diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index 99efbb5c0a..d2381f0276 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -711,6 +711,8 @@ container.bind(LOGS_SERVICE).toDynamicValue((ctx) => { }, readLocalLogs: (taskRunId: string) => ws.localLogs.read.query({ taskRunId }), + readLocalLogsTail: (taskRunId: string, maxBytes: number) => + ws.localLogs.readTail.query({ taskRunId, maxBytes }), writeLocalLogs: (taskRunId: string, content: string) => ws.localLogs.write.mutate({ taskRunId, content }), }; diff --git a/packages/core/src/sessions/sessionEventBatching.test.ts b/packages/core/src/sessions/sessionEventBatching.test.ts new file mode 100644 index 0000000000..93c436edfd --- /dev/null +++ b/packages/core/src/sessions/sessionEventBatching.test.ts @@ -0,0 +1,167 @@ +import type { AcpMessage, AgentSession } from "@posthog/shared"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { SessionService, type SessionServiceDeps } from "./sessionService"; + +const TASK_ID = "task-1"; +const RUN_ID = "run-1"; +const FLUSH_MS = 16; + +/** A plain streamed agent-message chunk — the common per-token event that just + * gets appended to the transcript. */ +function chunk(text: string): AcpMessage { + return { + ts: 1, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: RUN_ID, + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text }, + }, + }, + }, + } as unknown as AcpMessage; +} + +function chunkText(event: AcpMessage): string { + const params = (event.message as { params?: unknown }).params as { + update: { content: { text: string } }; + }; + return params.update.content.text; +} + +function createHarness() { + const sessions: Record = { + [RUN_ID]: { + taskRunId: RUN_ID, + taskId: TASK_ID, + events: [], + messageQueue: [], + pendingPermissions: new Map(), + status: "connected", + } as unknown as AgentSession, + }; + + const appendEvents = vi.fn( + (taskRunId: string, events: AcpMessage[], newLineCount?: number) => { + const session = sessions[taskRunId]; + if (!session) return; + session.events = [...session.events, ...events]; + if (newLineCount !== undefined) session.processedLineCount = newLineCount; + }, + ); + + const store = { + getSessions: () => sessions, + getSessionByTaskId: (taskId: string) => + Object.values(sessions).find((s) => s.taskId === taskId), + setSession: (session: AgentSession) => { + sessions[session.taskRunId] = session; + }, + updateSession: (taskRunId: string, updates: Partial) => { + const session = sessions[taskRunId]; + if (session) Object.assign(session, updates); + }, + appendEvents, + replaceOptimisticWithEvent: vi.fn(), + setPendingPermissions: vi.fn(), + clearMessageQueue: vi.fn(), + clearTailOptimisticItems: vi.fn(), + appendOptimisticItem: vi.fn(), + }; + + let onEvent: ((payload: unknown) => void) | undefined; + const noopLog = { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + }; + + const deps = { + store, + log: noopLog, + notifyPromptComplete: vi.fn(), + notifyPermissionRequest: vi.fn(), + taskViewedApi: { markActivity: vi.fn() }, + getPersistedConfigOptions: () => undefined, + setPersistedConfigOptions: vi.fn(), + trpc: { + agent: { + onSessionEvent: { + subscribe: ( + _input: unknown, + handlers: { onData: (payload: unknown) => void }, + ) => { + onEvent = handlers.onData; + return { unsubscribe: vi.fn() }; + }, + }, + onPermissionRequest: { + subscribe: () => ({ unsubscribe: vi.fn() }), + }, + onSessionIdleKilled: { + subscribe: () => ({ unsubscribe: vi.fn() }), + }, + }, + }, + } as unknown as SessionServiceDeps; + + const service = new SessionService(deps); + // Register the streamed-event subscription (captures onData). + ( + service as unknown as { subscribeToChannel(id: string): void } + ).subscribeToChannel(RUN_ID); + if (!onEvent) + throw new Error("subscribeToChannel did not subscribe to events"); + + return { + service, + appendEvents, + emit: (event: AcpMessage) => onEvent?.(event), + events: () => sessions[RUN_ID].events, + }; +} + +describe("streamed event batching", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + afterEach(() => { + vi.useRealTimers(); + }); + + it("defers a burst and applies it on one flush tick, in order", () => { + const h = createHarness(); + + h.emit(chunk("a")); + h.emit(chunk("b")); + h.emit(chunk("c")); + + // Nothing is applied synchronously — the burst is buffered. + expect(h.appendEvents).not.toHaveBeenCalled(); + expect(h.events()).toHaveLength(0); + + // A single flush tick drains the whole burst, in arrival order. + vi.advanceTimersByTime(FLUSH_MS); + expect(h.events().map(chunkText)).toEqual(["a", "b", "c"]); + }); + + it("flushes buffered events synchronously on teardown", () => { + const h = createHarness(); + + h.emit(chunk("a")); + h.emit(chunk("b")); + expect(h.events()).toHaveLength(0); + + // reset() tears down subscriptions and must not drop buffered events. + h.service.reset(); + expect(h.events().map(chunkText)).toEqual(["a", "b"]); + + // The flush timer was cleared, so advancing does not re-apply anything. + vi.advanceTimersByTime(FLUSH_MS); + expect(h.events()).toHaveLength(2); + }); +}); diff --git a/packages/core/src/sessions/sessionEventResidency.test.ts b/packages/core/src/sessions/sessionEventResidency.test.ts new file mode 100644 index 0000000000..2c2413562b --- /dev/null +++ b/packages/core/src/sessions/sessionEventResidency.test.ts @@ -0,0 +1,141 @@ +import type { AgentSession, SessionStatus } from "@posthog/shared"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { SessionService, type SessionServiceDeps } from "./sessionService"; +import { sessionStore, sessionStoreSetters } from "./sessionStore"; + +const RUN = "run-res"; +const TASK = "task-res"; +const GRACE_MS = 20_000; + +const LOG_LINE = JSON.stringify({ + type: "notification", + notification: { + method: "session/update", + params: { + sessionId: RUN, + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: "restored" }, + }, + }, + }, +}); + +function makeService(readLocalLogs = vi.fn().mockResolvedValue("")) { + const deps = { + store: sessionStoreSetters, + log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }, + notifyPromptComplete: vi.fn(), + notifyPermissionRequest: vi.fn(), + taskViewedApi: { markActivity: vi.fn() }, + getPersistedConfigOptions: () => undefined, + setPersistedConfigOptions: vi.fn(), + trpc: { + agent: { + onSessionEvent: { subscribe: () => ({ unsubscribe: vi.fn() }) }, + onPermissionRequest: { subscribe: () => ({ unsubscribe: vi.fn() }) }, + onSessionIdleKilled: { subscribe: () => ({ unsubscribe: vi.fn() }) }, + }, + logs: { readLocalLogs: { query: readLocalLogs } }, + }, + } as unknown as SessionServiceDeps; + return new SessionService(deps); +} + +function seed(status: SessionStatus, isPromptPending = false) { + sessionStoreSetters.setSession({ + taskRunId: RUN, + taskId: TASK, + events: [], + messageQueue: [], + pendingPermissions: new Map(), + status, + isPromptPending, + } as unknown as AgentSession); + sessionStoreSetters.appendEvents(RUN, [{ ts: 1, message: {} } as never]); +} + +const events = () => sessionStore.getState().sessions[RUN]?.events ?? []; + +describe("session transcript residency", () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => { + vi.useRealTimers(); + sessionStoreSetters.removeSession(RUN); + }); + + it("evicts a disconnected, idle session after the grace window", () => { + const service = makeService(); + seed("disconnected"); + service.scheduleEventEviction(TASK); + + expect(events()).toHaveLength(1); + vi.advanceTimersByTime(GRACE_MS); + expect(events()).toHaveLength(0); + }); + + it("never evicts a connected session", () => { + const service = makeService(); + seed("connected"); + service.scheduleEventEviction(TASK); + + vi.advanceTimersByTime(GRACE_MS); + expect(events()).toHaveLength(1); + }); + + it("never evicts a session with a prompt in flight", () => { + const service = makeService(); + seed("disconnected", true); + service.scheduleEventEviction(TASK); + + vi.advanceTimersByTime(GRACE_MS); + expect(events()).toHaveLength(1); + }); + + it("ensureEventsLoaded cancels a pending eviction", () => { + const service = makeService(); + seed("disconnected"); + service.scheduleEventEviction(TASK); + + void service.ensureEventsLoaded(TASK); // return to the view before grace + vi.advanceTimersByTime(GRACE_MS); + expect(events()).toHaveLength(1); + }); + + it("rehydrates an evicted transcript from disk on return", async () => { + const readLocalLogs = vi.fn().mockResolvedValue(LOG_LINE); + const service = makeService(readLocalLogs); + seed("disconnected"); + + service.scheduleEventEviction(TASK); + vi.advanceTimersByTime(GRACE_MS); + expect(events()).toHaveLength(0); + + await service.ensureEventsLoaded(TASK); + expect(readLocalLogs).toHaveBeenCalledWith({ taskRunId: RUN }); + expect(events()).toHaveLength(1); + }); + + it("retries rehydration after a failed log read instead of stranding it", async () => { + const readLocalLogs = vi + .fn() + .mockRejectedValueOnce(new Error("transient")) + .mockResolvedValueOnce(LOG_LINE); + const service = makeService(readLocalLogs); + seed("disconnected"); + + service.scheduleEventEviction(TASK); + vi.advanceTimersByTime(GRACE_MS); + expect(events()).toHaveLength(0); + + // First visit: the log read throws — the transcript stays empty but the + // run is re-marked evicted so a later visit can retry. + await service.ensureEventsLoaded(TASK); + expect(events()).toHaveLength(0); + + // Second visit: the read succeeds and the transcript is restored. + await service.ensureEventsLoaded(TASK); + expect(events()).toHaveLength(1); + expect(readLocalLogs).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/core/src/sessions/sessionEvents.test.ts b/packages/core/src/sessions/sessionEvents.test.ts index aa5771b442..332f630179 100644 --- a/packages/core/src/sessions/sessionEvents.test.ts +++ b/packages/core/src/sessions/sessionEvents.test.ts @@ -237,6 +237,14 @@ describe("convertStoredEntriesToEvents — imported user prompts", () => { const msg = events[0].message; expect("method" in msg && msg.method).toBe("session/update"); }); + + it("freezes converted events on both the promoted and raw branches", () => { + const events = convertStoredEntriesToEvents([ + userChunkEntry("promoted", { importedUserPrompt: true }), + userChunkEntry("raw"), + ]); + expect(events.every((event) => Object.isFrozen(event))).toBe(true); + }); }); describe("isAbsoluteFolderPath", () => { diff --git a/packages/core/src/sessions/sessionEvents.ts b/packages/core/src/sessions/sessionEvents.ts index a8fccff790..aa6cd24148 100644 --- a/packages/core/src/sessions/sessionEvents.ts +++ b/packages/core/src/sessions/sessionEvents.ts @@ -29,12 +29,14 @@ import { extractPromptDisplayContent } from "./promptContent"; function storedEntryToAcpMessage(entry: StoredLogEntry): AcpMessage { const ts = entry.timestamp ? new Date(entry.timestamp).getTime() : Date.now(); const promoted = promoteImportedUserPrompt(entry, ts); - if (promoted) return promoted; - return { + // Freeze at creation: events assigned via setSession bypass the store's + // per-append freeze, so this keeps them read-only once stored. + if (promoted) return Object.freeze(promoted); + return Object.freeze({ type: "acp_message", ts, message: (entry.notification ?? {}) as JsonRpcMessage, - }; + }); } /** diff --git a/packages/core/src/sessions/sessionOpenTailFirst.test.ts b/packages/core/src/sessions/sessionOpenTailFirst.test.ts new file mode 100644 index 0000000000..9cabaf6e86 --- /dev/null +++ b/packages/core/src/sessions/sessionOpenTailFirst.test.ts @@ -0,0 +1,105 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { SessionService, type SessionServiceDeps } from "./sessionService"; +import { sessionStore, sessionStoreSetters } from "./sessionStore"; + +const RUN = "run-tf"; +const TASK = "task-tf"; + +function line(text: string): string { + return JSON.stringify({ + type: "notification", + notification: { + method: "session/update", + params: { + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text }, + }, + }, + }, + }); +} + +function makeService(readTail?: unknown) { + const logs: Record = { + readLocalLogs: { query: vi.fn().mockResolvedValue(null) }, + fetchS3Logs: { query: vi.fn().mockResolvedValue(null) }, + writeLocalLogs: { mutate: vi.fn() }, + }; + if (readTail !== undefined) logs.readLocalLogsTail = { query: readTail }; + + const deps = { + store: sessionStoreSetters, + log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }, + notifyPromptComplete: vi.fn(), + notifyPermissionRequest: vi.fn(), + taskViewedApi: { markActivity: vi.fn() }, + getPersistedConfigOptions: () => undefined, + setPersistedConfigOptions: vi.fn(), + trpc: { + agent: { + onSessionEvent: { subscribe: () => ({ unsubscribe: vi.fn() }) }, + onPermissionRequest: { subscribe: () => ({ unsubscribe: vi.fn() }) }, + onSessionIdleKilled: { subscribe: () => ({ unsubscribe: vi.fn() }) }, + }, + logs, + }, + } as unknown as SessionServiceDeps; + return new SessionService(deps); +} + +type Painter = { + paintTailFirst(r: string, t: string, ti: string, u: string): Promise; +}; +const paint = (svc: SessionService) => + (svc as unknown as Painter).paintTailFirst(RUN, TASK, "Title", "log-url"); + +const events = () => sessionStore.getState().sessions[RUN]?.events ?? []; + +afterEach(() => sessionStoreSetters.removeSession(RUN)); + +describe("paintTailFirst", () => { + it("paints a session from the tail content", async () => { + const readTail = vi.fn().mockResolvedValue({ + content: `${line("a")}\n${line("b")}\n`, + truncated: true, + }); + await paint(makeService(readTail)); + + expect(readTail).toHaveBeenCalledWith({ + taskRunId: RUN, + maxBytes: 1_500_000, + }); + expect(events().length).toBeGreaterThan(0); + expect(sessionStore.getState().sessions[RUN]?.logUrl).toBe("log-url"); + }); + + it("is a no-op when the host doesn't expose the tail read", async () => { + await paint(makeService(undefined)); + expect(sessionStore.getState().sessions[RUN]).toBeUndefined(); + }); + + it("is a no-op when a session already exists", async () => { + const readTail = vi + .fn() + .mockResolvedValue({ content: line("x"), truncated: true }); + sessionStoreSetters.setSession({ + taskRunId: RUN, + taskId: TASK, + events: [], + messageQueue: [], + pendingPermissions: new Map(), + status: "connected", + } as never); + await paint(makeService(readTail)); + expect(readTail).not.toHaveBeenCalled(); + }); + + it("is a no-op on empty tail content", async () => { + const readTail = vi + .fn() + .mockResolvedValue({ content: " ", truncated: true }); + await paint(makeService(readTail)); + expect(sessionStore.getState().sessions[RUN]).toBeUndefined(); + }); +}); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 6dd03c0d18..5b5af8740e 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -94,6 +94,28 @@ const AUTO_RETRY_MAX_ATTEMPTS = 2; const AUTO_RETRY_DELAY_MS = 10_000; const AUTH_RESTORE_MAX_RETRY_WAITS = 6; const MAX_SUPERSEDED_RUN_IDS = 100; +/** + * Streamed events are buffered and flushed on this cadence so a burst of tokens + * coalesces into one processing pass (and roughly one render) instead of one + * per event. Electron IPC delivers each event as its own task, so a microtask + * flush wouldn't batch across them — a short timer does. One frame is + * imperceptible for streamed text. + */ +const SESSION_EVENT_FLUSH_MS = 16; +/** + * A backgrounded session's transcript is freed this long after it stops being + * viewed, and reloaded from disk on return. Only disconnected (idle, no live + * subscription) sessions are eligible, so no streamed event can append to an + * evicted transcript. + */ +const SESSION_EVENT_EVICT_GRACE_MS = 20_000; +/** + * On open, paint the last this-many bytes of the log immediately so a big + * transcript shows its latest turns in tens of ms, while the authoritative + * full read + connect completes behind it. ~1.5MB is a few hundred entries — + * plenty for the initial (scrolled-to-bottom) view. + */ +const OPEN_TAIL_BYTES = 1_500_000; class GitHubAuthorizationRequiredForCloudHandoffError extends Error { constructor( @@ -146,6 +168,9 @@ export interface SessionTrpc { }; logs: { readLocalLogs: TrpcQuery; + /** Optional: only the Electron host exposes the tail read. Core feature- + * detects and falls back to a full read when it's absent. */ + readLocalLogsTail?: TrpcQuery; fetchS3Logs: TrpcQuery; writeLocalLogs: TrpcMutation; }; @@ -161,6 +186,12 @@ export interface ISessionStore { events: AcpMessage[], newLineCount?: number, ): void; + evictEvents(taskRunId: string): void; + restoreEvents( + taskRunId: string, + events: AcpMessage[], + lineCount: number, + ): void; updateCloudStatus( taskRunId: string, fields: { @@ -672,9 +703,13 @@ export class SessionService { return; } + // Paint the log tail immediately so a big transcript is visible in tens + // of ms; the full read + reconnect replace it with the authoritative + // session once everything below resolves. const [workspaceResult, logResult] = await Promise.all([ this.d.trpc.workspace.verify.query({ taskId }), this.fetchSessionLogs(logUrl, existingRunId), + this.paintTailFirst(existingRunId, taskId, taskTitle, logUrl), ]); if (!workspaceResult.exists) { @@ -1010,6 +1045,8 @@ export class SessionService { } this.unsubscribeFromChannel(taskRunId); + this.cancelEventEviction(taskRunId); + this.evictedRunIds.delete(taskRunId); this.d.store.removeSession(taskRunId); this.cloudRunIdleTracker.delete(taskRunId); this.cloudLogGapReconciler.forgetDeficiency(taskRunId); @@ -1323,6 +1360,142 @@ export class SessionService { // --- Subscription Management --- + /** Streamed events awaiting their frame flush, keyed by taskRunId. Order + * within a taskRunId is preserved; taskRunIds are independent. */ + private pendingSessionEvents = new Map(); + private sessionEventFlushHandle: ReturnType | null = null; + + private enqueueSessionEvent(taskRunId: string, acpMsg: AcpMessage): void { + const buffered = this.pendingSessionEvents.get(taskRunId); + if (buffered) { + buffered.push(acpMsg); + } else { + this.pendingSessionEvents.set(taskRunId, [acpMsg]); + } + if (this.sessionEventFlushHandle === null) { + this.sessionEventFlushHandle = setTimeout(() => { + this.sessionEventFlushHandle = null; + this.flushSessionEvents(); + }, SESSION_EVENT_FLUSH_MS); + } + } + + private flushSessionEvents(): void { + if (this.pendingSessionEvents.size === 0) return; + const batches = this.pendingSessionEvents; + this.pendingSessionEvents = new Map(); + for (const [taskRunId, events] of batches) { + for (const acpMsg of events) { + this.handleSessionEvent(taskRunId, acpMsg); + } + } + } + + /** Drain one task's buffer immediately, so a reader (permission handling, + * teardown) never sees a transcript missing already-received events. */ + private flushSessionEventsForTask(taskRunId: string): void { + const events = this.pendingSessionEvents.get(taskRunId); + if (!events) return; + this.pendingSessionEvents.delete(taskRunId); + for (const acpMsg of events) { + this.handleSessionEvent(taskRunId, acpMsg); + } + } + + // --- Transcript residency (memory eviction) --- + + /** taskRunIds whose transcript was freed and must be reloaded on next view. */ + private evictedRunIds = new Set(); + private eventEvictionTimers = new Map< + string, + ReturnType + >(); + + /** + * Called when a task's transcript becomes visible. Cancels any pending + * eviction and, if the transcript was freed while backgrounded, reloads it + * from disk — but only if a reconnect hasn't already refilled it. + */ + async ensureEventsLoaded(taskId: string): Promise { + const session = this.d.store.getSessionByTaskId(taskId); + if (!session) return; + const { taskRunId } = session; + this.cancelEventEviction(taskRunId); + if (!this.evictedRunIds.has(taskRunId)) return; + + try { + if (session.events.length === 0) { + const { rawEntries, totalLineCount } = await this.fetchSessionLogs( + session.logUrl, + taskRunId, + ); + // A reconnect may have refilled events while we awaited the log read; + // only restore if the transcript is still empty for the same run. + const fresh = this.d.store.getSessionByTaskId(taskId); + if ( + fresh?.taskRunId === taskRunId && + fresh.events.length === 0 && + rawEntries.length > 0 + ) { + this.d.store.restoreEvents( + taskRunId, + convertStoredEntriesToEvents(rawEntries), + totalLineCount, + ); + } + } + // Clear the evicted flag only once the transcript is populated — restored + // here, or refilled by a reconnect. An empty read leaves the run evicted so + // a later visit retries: fetchSessionLogs swallows read errors and returns + // empty rather than throwing, so a transient failure would otherwise strand + // the transcript empty permanently. + if ((this.d.store.getSessionByTaskId(taskId)?.events.length ?? 0) > 0) { + this.evictedRunIds.delete(taskRunId); + } + } catch (error) { + this.d.log.warn("Failed to rehydrate evicted session transcript", { + taskId, + error, + }); + } + } + + /** + * Called when a task's transcript stops being visible. Schedules its + * transcript to be freed after a grace period, if it's still a settled, + * disconnected background session by then. + */ + scheduleEventEviction(taskId: string): void { + const session = this.d.store.getSessionByTaskId(taskId); + if (!session) return; + const { taskRunId } = session; + if (this.eventEvictionTimers.has(taskRunId)) return; + + const timer = setTimeout(() => { + this.eventEvictionTimers.delete(taskRunId); + const current = this.d.store.getSessions()[taskRunId]; + if ( + !current || + current.status !== "disconnected" || + current.isPromptPending || + current.events.length === 0 + ) { + return; + } + this.evictedRunIds.add(taskRunId); + this.d.store.evictEvents(taskRunId); + }, SESSION_EVENT_EVICT_GRACE_MS); + this.eventEvictionTimers.set(taskRunId, timer); + } + + private cancelEventEviction(taskRunId: string): void { + const timer = this.eventEvictionTimers.get(taskRunId); + if (timer !== undefined) { + clearTimeout(timer); + this.eventEvictionTimers.delete(taskRunId); + } + } + private subscribeToChannel(taskRunId: string): void { if (this.subscriptions.has(taskRunId)) { return; @@ -1332,7 +1505,7 @@ export class SessionService { { taskRunId }, { onData: (payload: unknown) => { - this.handleSessionEvent(taskRunId, payload as AcpMessage); + this.enqueueSessionEvent(taskRunId, payload as AcpMessage); }, onError: (err) => { this.d.log.error("Session subscription error", { @@ -1383,6 +1556,9 @@ export class SessionService { } private unsubscribeFromChannel(taskRunId: string): void { + // Apply anything still buffered before we stop listening, so a closing + // channel doesn't drop its final events. + this.flushSessionEventsForTask(taskRunId); const subscription = this.subscriptions.get(taskRunId); subscription?.event.unsubscribe(); subscription?.permission?.unsubscribe(); @@ -1411,6 +1587,14 @@ export class SessionService { this.stopCloudTaskWatch(taskId); } + if (this.sessionEventFlushHandle !== null) { + clearTimeout(this.sessionEventFlushHandle); + this.sessionEventFlushHandle = null; + } + this.pendingSessionEvents.clear(); + for (const timer of this.eventEvictionTimers.values()) clearTimeout(timer); + this.eventEvictionTimers.clear(); + this.evictedRunIds.clear(); this.connectingTasks.clear(); this.localRepoPaths.clear(); this.localRecoveryAttempts.clear(); @@ -1801,6 +1985,10 @@ export class SessionService { title: payload.toolCall.title, }); + // A permission request references a tool call from the stream; apply any + // buffered events first so that tool call is present in the transcript. + this.flushSessionEventsForTask(taskRunId); + // Get fresh session state const session = this.d.store.getSessions()[taskRunId]; if (!session) { @@ -4620,6 +4808,43 @@ export class SessionService { }); } + /** + * Paint the tail of a task's local log immediately so a big transcript shows + * its latest turns in tens of ms, instead of blocking on the full-log read + + * IPC transfer. This is a throwaway fast-paint: the authoritative full read + + * connect (`reconnectToLocalSession`) replaces this session shortly after with + * correct processed-line tracking. No-op when a session already exists, the + * host doesn't expose the tail read, or there's no local log. + */ + private async paintTailFirst( + taskRunId: string, + taskId: string, + taskTitle: string, + logUrl: string, + ): Promise { + const tailQuery = this.d.trpc.logs.readLocalLogsTail; + if (!tailQuery) return; + if (this.d.store.getSessionByTaskId(taskId)) return; + try { + const res = (await tailQuery.query({ + taskRunId, + maxBytes: OPEN_TAIL_BYTES, + })) as { content: string; truncated: boolean } | null; + if (!res?.content?.trim()) return; + // The full read may have set the session while we awaited the tail. + if (this.d.store.getSessionByTaskId(taskId)) return; + const { rawEntries } = this.parseLogContent(res.content); + if (rawEntries.length === 0) return; + const session = createBaseSession(taskRunId, taskId, taskTitle); + session.events = convertStoredEntriesToEvents(rawEntries); + session.logUrl = logUrl; + session.status = "connecting"; + this.d.store.setSession(session); + } catch (error) { + this.d.log.debug("Tail-first paint skipped", { taskId, error }); + } + } + private async fetchSessionLogs( logUrl: string | undefined, taskRunId?: string, diff --git a/packages/core/src/sessions/sessionStore.ts b/packages/core/src/sessions/sessionStore.ts index 2ea1cbad09..0e86a78c6a 100644 --- a/packages/core/src/sessions/sessionStore.ts +++ b/packages/core/src/sessions/sessionStore.ts @@ -7,9 +7,17 @@ import type { QueuedMessage, TaskRunStatus, } from "@posthog/shared"; +import { setAutoFreeze } from "immer"; import { immer } from "zustand/middleware/immer"; import { createStore } from "zustand/vanilla"; +// immer autofreeze deep-walks produced state on every commit. For the +// append-only `events` array that re-walks the whole (growing) array on every +// streamed event — O(n) per append, O(n²) per turn. Autofreeze is a dev-time +// mutation guard with no runtime value, so disable it; events are frozen +// individually at the append/creation seam instead, which is O(1) each. +setAutoFreeze(false); + export interface SessionState { /** Sessions indexed by taskRunId */ sessions: Record; @@ -64,6 +72,9 @@ export const sessionStoreSetters = { sessionStore.setState((state) => { const session = state.sessions[taskRunId]; if (session) { + // Keep each event immutable once stored (O(1) each). The store disables + // immer autofreeze, so this is the only freeze. + for (const event of events) Object.freeze(event); session.events.push(...events); if (newLineCount !== undefined) { session.processedLineCount = newLineCount; @@ -72,6 +83,40 @@ export const sessionStoreSetters = { }); }, + /** + * Free a backgrounded session's transcript to reclaim memory. The events are + * reloaded from disk the next time the session is viewed (see + * `SessionService.ensureEventsLoaded`). No-op if the session is gone. + */ + evictEvents: (taskRunId: string) => { + sessionStore.setState((state) => { + const session = state.sessions[taskRunId]; + if (session && session.events.length > 0) { + session.events = []; + session.processedLineCount = 0; + } + }); + }, + + /** + * Replace a session's transcript in place (rehydration after eviction), + * preserving its live status/config. No-op if the session is gone. + */ + restoreEvents: ( + taskRunId: string, + events: AcpMessage[], + lineCount: number, + ) => { + sessionStore.setState((state) => { + const session = state.sessions[taskRunId]; + if (session) { + for (const event of events) Object.freeze(event); + session.events = events; + session.processedLineCount = lineCount; + } + }); + }, + updateCloudStatus: ( taskRunId: string, fields: { @@ -255,7 +300,7 @@ export const sessionStoreSetters = { sessionStore.setState((state) => { const session = state.sessions[taskRunId]; if (session) { - session.events.push(event); + session.events.push(Object.freeze(event)); session.optimisticItems = []; } }); diff --git a/packages/core/src/sessions/sessionStoreEviction.test.ts b/packages/core/src/sessions/sessionStoreEviction.test.ts new file mode 100644 index 0000000000..da27d09f9c --- /dev/null +++ b/packages/core/src/sessions/sessionStoreEviction.test.ts @@ -0,0 +1,79 @@ +import type { AcpMessage, AgentSession } from "@posthog/shared"; +import { afterEach, describe, expect, it } from "vitest"; +import { sessionStore, sessionStoreSetters } from "./sessionStore"; + +const RUN = "run-evict"; +const TASK = "task-evict"; + +function seedWithEvents() { + sessionStoreSetters.setSession({ + taskRunId: RUN, + taskId: TASK, + events: [], + messageQueue: [], + pendingPermissions: new Map(), + status: "disconnected", + } as unknown as AgentSession); + sessionStoreSetters.appendEvents( + RUN, + [{ ts: 1, message: {} } as unknown as AcpMessage], + 3, + ); +} + +afterEach(() => sessionStoreSetters.removeSession(RUN)); + +describe("evictEvents / restoreEvents", () => { + it("evictEvents frees the transcript and resets the line cursor", () => { + seedWithEvents(); + expect(sessionStore.getState().sessions[RUN].events).toHaveLength(1); + + sessionStoreSetters.evictEvents(RUN); + + const s = sessionStore.getState().sessions[RUN]; + expect(s.events).toHaveLength(0); + expect(s.processedLineCount).toBe(0); + }); + + it("restoreEvents refills the transcript and freezes each event", () => { + seedWithEvents(); + sessionStoreSetters.evictEvents(RUN); + + sessionStoreSetters.restoreEvents( + RUN, + [{ ts: 2, message: {} } as unknown as AcpMessage], + 7, + ); + + const s = sessionStore.getState().sessions[RUN]; + expect(s.events).toHaveLength(1); + expect(s.processedLineCount).toBe(7); + expect(Object.isFrozen(s.events[0])).toBe(true); + }); + + it("appendEvents and replaceOptimisticWithEvent freeze each stored event", () => { + seedWithEvents(); + sessionStoreSetters.replaceOptimisticWithEvent(RUN, { + ts: 2, + message: {}, + } as unknown as AcpMessage); + + const s = sessionStore.getState().sessions[RUN]; + expect(s.events).toHaveLength(2); + expect(s.events.every((event) => Object.isFrozen(event))).toBe(true); + }); + + it("evictEvents is a no-op on an already-empty session", () => { + sessionStoreSetters.setSession({ + taskRunId: RUN, + taskId: TASK, + events: [], + messageQueue: [], + pendingPermissions: new Map(), + status: "disconnected", + } as unknown as AgentSession); + + expect(() => sessionStoreSetters.evictEvents(RUN)).not.toThrow(); + expect(sessionStore.getState().sessions[RUN].events).toHaveLength(0); + }); +}); diff --git a/packages/core/src/sidebar/buildSidebarData.ts b/packages/core/src/sidebar/buildSidebarData.ts index 5b41353260..43edfc52b2 100644 --- a/packages/core/src/sidebar/buildSidebarData.ts +++ b/packages/core/src/sidebar/buildSidebarData.ts @@ -88,6 +88,29 @@ export interface TaskSession { cloudOutput?: { pr_url?: unknown } | null; } +/** + * A primitive signature of just the session fields the sidebar renders (see + * {@link deriveTaskData}). The sidebar subscribes to this instead of the whole + * sessions record, so it doesn't rebuild on every streamed event — only when a + * field it actually reads changes. It deliberately ignores `events`. + */ +export function computeSidebarSessionSignature( + sessions: Record, +): string { + let signature = ""; + for (const session of Object.values(sessions)) { + if (!session.taskId) continue; + const prUrl = + typeof session.cloudOutput?.pr_url === "string" + ? session.cloudOutput.pr_url + : ""; + signature += `${session.taskId}:${session.isPromptPending ? 1 : 0}:${ + session.pendingPermissions?.size ?? 0 + }:${session.cloudStatus ?? ""}:${prUrl};`; + } + return signature; +} + export interface TaskWorkspace { folderId?: string | null; folderPath?: string | null; diff --git a/packages/core/src/sidebar/computeSidebarSessionSignature.test.ts b/packages/core/src/sidebar/computeSidebarSessionSignature.test.ts new file mode 100644 index 0000000000..4b12cdc0b6 --- /dev/null +++ b/packages/core/src/sidebar/computeSidebarSessionSignature.test.ts @@ -0,0 +1,49 @@ +import { describe, expect, it } from "vitest"; +import { + computeSidebarSessionSignature, + type TaskSession, +} from "./buildSidebarData"; + +type SigInput = Record; + +function sig(sessions: Record): string { + return computeSidebarSessionSignature(sessions as SigInput); +} + +describe("computeSidebarSessionSignature", () => { + it("ignores fields the sidebar doesn't render (e.g. events)", () => { + const before = sig({ + r1: { taskId: "t1", isPromptPending: true, events: [1] }, + }); + const after = sig({ + r1: { taskId: "t1", isPromptPending: true, events: [1, 2, 3, 4] }, + }); + expect(after).toBe(before); + }); + + it("changes when isPromptPending flips", () => { + const a = sig({ r1: { taskId: "t1", isPromptPending: false } }); + const b = sig({ r1: { taskId: "t1", isPromptPending: true } }); + expect(a).not.toBe(b); + }); + + it("changes when the pending-permission count changes", () => { + const a = sig({ r1: { taskId: "t1", pendingPermissions: { size: 0 } } }); + const b = sig({ r1: { taskId: "t1", pendingPermissions: { size: 1 } } }); + expect(a).not.toBe(b); + }); + + it("changes when cloud status or PR url changes", () => { + const a = sig({ r1: { taskId: "t1", cloudStatus: "running" } }); + const b = sig({ r1: { taskId: "t1", cloudStatus: "completed" } }); + expect(a).not.toBe(b); + + const c = sig({ r1: { taskId: "t1", cloudOutput: { pr_url: "x" } } }); + const d = sig({ r1: { taskId: "t1", cloudOutput: { pr_url: "y" } } }); + expect(c).not.toBe(d); + }); + + it("skips sessions without a taskId", () => { + expect(sig({ r1: { isPromptPending: true } })).toBe(""); + }); +}); diff --git a/packages/host-router/src/routers/logs.router.ts b/packages/host-router/src/routers/logs.router.ts index f2efdb2245..2c4e177be4 100644 --- a/packages/host-router/src/routers/logs.router.ts +++ b/packages/host-router/src/routers/logs.router.ts @@ -6,6 +6,8 @@ import { fetchS3LogsOutput, readLocalLogsInput, readLocalLogsOutput, + readLocalLogsTailInput, + readLocalLogsTailOutput, writeLocalLogsInput, } from "@posthog/workspace-server/services/local-logs/schemas"; @@ -26,6 +28,15 @@ export const logsRouter = router({ .readLocalLogs(input.taskRunId), ), + readLocalLogsTail: publicProcedure + .input(readLocalLogsTailInput) + .output(readLocalLogsTailOutput) + .query(({ ctx, input }) => + ctx.container + .get(LOGS_SERVICE) + .readLocalLogsTail(input.taskRunId, input.maxBytes), + ), + writeLocalLogs: publicProcedure .input(writeLocalLogsInput) .mutation(({ ctx, input }) => diff --git a/packages/ui/src/features/code-review/components/LazyReviewPages.tsx b/packages/ui/src/features/code-review/components/LazyReviewPages.tsx new file mode 100644 index 0000000000..f261f0ab89 --- /dev/null +++ b/packages/ui/src/features/code-review/components/LazyReviewPages.tsx @@ -0,0 +1,37 @@ +import type { Task } from "@posthog/shared/domain-types"; +import { DotsCircleSpinner } from "@posthog/ui/primitives/DotsCircleSpinner"; +import { lazy, type ReactNode, Suspense } from "react"; + +// The code-review surface (ReviewShell, diff rows, comment UI, review hooks) is +// only reached when a review is opened, so it's split out of the initial bundle. +// The underlying diff/highlight libraries stay eager — the transcript uses them. +const ReviewPageLazy = lazy(() => + import("./ReviewPage").then((m) => ({ default: m.ReviewPage })), +); +const CloudReviewPageLazy = lazy(() => + import("./CloudReviewPage").then((m) => ({ default: m.CloudReviewPage })), +); + +function ReviewFallback(): ReactNode { + return ( +
+ +
+ ); +} + +export function LazyReviewPage({ task }: { task: Task }): ReactNode { + return ( + }> + + + ); +} + +export function LazyCloudReviewPage({ task }: { task: Task }): ReactNode { + return ( + }> + + + ); +} diff --git a/packages/ui/src/features/editor/components/StreamingMarkdown.tsx b/packages/ui/src/features/editor/components/StreamingMarkdown.tsx index 0293e71cb3..77fff28afe 100644 --- a/packages/ui/src/features/editor/components/StreamingMarkdown.tsx +++ b/packages/ui/src/features/editor/components/StreamingMarkdown.tsx @@ -1,5 +1,5 @@ import { CodeBlock } from "@posthog/ui/primitives/CodeBlock"; -import { memo } from "react"; +import { memo, useMemo } from "react"; import type { Components } from "react-markdown"; import { MarkdownRenderer } from "./MarkdownRenderer"; import { parseOpenFence, splitMarkdownBlocks } from "./splitMarkdownBlocks"; @@ -26,7 +26,7 @@ export const StreamingMarkdown = memo(function StreamingMarkdown({ content, componentsOverride, }: StreamingMarkdownProps) { - const blocks = splitMarkdownBlocks(content); + const blocks = useMemo(() => splitMarkdownBlocks(content), [content]); const lastIndex = blocks.length - 1; return ( diff --git a/packages/ui/src/features/panels/createDebouncedStorage.test.ts b/packages/ui/src/features/panels/createDebouncedStorage.test.ts new file mode 100644 index 0000000000..09ddd9f515 --- /dev/null +++ b/packages/ui/src/features/panels/createDebouncedStorage.test.ts @@ -0,0 +1,61 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { createDebouncedStorage } from "./panelLayoutStore"; + +function fakeBase() { + return { + getItem: vi.fn(() => null as string | null), + setItem: vi.fn(), + removeItem: vi.fn(), + }; +} + +describe("createDebouncedStorage", () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + it("coalesces rapid writes to the same key into a single write", () => { + const base = fakeBase(); + const storage = createDebouncedStorage(base, 200); + + storage.setItem("k", "a"); + storage.setItem("k", "b"); + storage.setItem("k", "c"); + expect(base.setItem).not.toHaveBeenCalled(); + + vi.advanceTimersByTime(200); + expect(base.setItem).toHaveBeenCalledTimes(1); + expect(base.setItem).toHaveBeenCalledWith("k", "c"); + }); + + it("passes reads through synchronously", () => { + const base = fakeBase(); + base.getItem.mockReturnValue("v"); + const storage = createDebouncedStorage(base, 200); + + expect(storage.getItem("k")).toBe("v"); + expect(base.getItem).toHaveBeenCalledWith("k"); + }); + + it("cancels a pending write when the key is removed", () => { + const base = fakeBase(); + const storage = createDebouncedStorage(base, 200); + + storage.setItem("k", "a"); + storage.removeItem("k"); + vi.advanceTimersByTime(200); + + expect(base.setItem).not.toHaveBeenCalled(); + expect(base.removeItem).toHaveBeenCalledWith("k"); + }); + + it("debounces different keys independently", () => { + const base = fakeBase(); + const storage = createDebouncedStorage(base, 200); + + storage.setItem("a", "1"); + storage.setItem("b", "2"); + vi.advanceTimersByTime(200); + + expect(base.setItem).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/ui/src/features/panels/panelLayoutStore.test.ts b/packages/ui/src/features/panels/panelLayoutStore.test.ts index 6239049b96..d005ff2cfd 100644 --- a/packages/ui/src/features/panels/panelLayoutStore.test.ts +++ b/packages/ui/src/features/panels/panelLayoutStore.test.ts @@ -286,6 +286,8 @@ describe("panelLayoutStore", () => { usePanelLayoutStore.getState().initializeTask("task-1"); usePanelLayoutStore.getState().openFile("task-1", "src/App.tsx"); + // Persistence is debounced; pagehide flushes pending writes. + window.dispatchEvent(new Event("pagehide")); const storedData = localStorage.getItem("panel-layout-store"); expect(storedData).not.toBeNull(); @@ -300,6 +302,8 @@ describe("panelLayoutStore", () => { usePanelLayoutStore.getState().initializeTask("task-1"); usePanelLayoutStore.getState().openFile("task-1", "src/App.tsx"); + // Persistence is debounced; pagehide flushes pending writes. + window.dispatchEvent(new Event("pagehide")); const storedData = localStorage.getItem("panel-layout-store"); usePanelLayoutStore.getState().clearAllLayouts(); diff --git a/packages/ui/src/features/panels/panelLayoutStore.ts b/packages/ui/src/features/panels/panelLayoutStore.ts index 14588a3869..12888543b6 100644 --- a/packages/ui/src/features/panels/panelLayoutStore.ts +++ b/packages/ui/src/features/panels/panelLayoutStore.ts @@ -21,7 +21,11 @@ import { import { createFileTabId } from "@posthog/core/panels/panelStoreHelpers"; import { findTabInTree } from "@posthog/core/panels/panelTree"; import { ANALYTICS_EVENTS, getFileExtension } from "@posthog/shared"; -import { persist } from "zustand/middleware"; +import { + createJSONStorage, + persist, + type StateStorage, +} from "zustand/middleware"; import { createWithEqualityFn } from "zustand/traditional"; import { track } from "../../shell/analytics"; import { updateTaskLayout } from "./panelStoreHelpers"; @@ -113,6 +117,72 @@ export interface PanelLayoutStore { clearAllLayouts: () => void; } +const PANEL_PERSIST_DEBOUNCE_MS = 200; + +/** + * Wraps a storage so writes to the same key coalesce onto a trailing debounce. + * Reads stay synchronous and in-memory state is untouched, so live UI is + * unaffected; only the write to the backing store is deferred. Pending writes + * flush on `pagehide` so the last change before the window closes isn't lost. + */ +export function createDebouncedStorage( + base: StateStorage, + waitMs: number, +): StateStorage { + const pending = new Map(); + const timers = new Map>(); + + const flush = (key: string) => { + timers.delete(key); + const value = pending.get(key); + pending.delete(key); + if (value !== undefined) base.setItem(key, value); + }; + + if (typeof window !== "undefined") { + window.addEventListener("pagehide", () => { + for (const key of [...pending.keys()]) flush(key); + }); + } + + return { + getItem: (key) => base.getItem(key), + setItem: (key, value) => { + pending.set(key, value); + const existing = timers.get(key); + if (existing !== undefined) clearTimeout(existing); + timers.set( + key, + setTimeout(() => flush(key), waitMs), + ); + }, + removeItem: (key) => { + const existing = timers.get(key); + if (existing !== undefined) { + clearTimeout(existing); + timers.delete(key); + } + pending.delete(key); + base.removeItem(key); + }, + }; +} + +/** + * react-resizable-panels fires a layout change every frame during a drag, and + * persist serializes the whole layout tree on each one. Panels are uncontrolled + * (defaultSize), so debouncing the write keeps live resize instant while + * collapsing a drag's ~60 synchronous localStorage writes into one. + */ +const panelLayoutStorage: StateStorage = createDebouncedStorage( + { + getItem: (key) => window.localStorage.getItem(key), + setItem: (key, value) => window.localStorage.setItem(key, value), + removeItem: (key) => window.localStorage.removeItem(key), + }, + PANEL_PERSIST_DEBOUNCE_MS, +); + export const usePanelLayoutStore = createWithEqualityFn()( persist( (set, get) => ({ @@ -414,6 +484,7 @@ export const usePanelLayoutStore = createWithEqualityFn()( name: "panel-layout-store", version: 10, migrate: () => ({ taskLayouts: {} }), + storage: createJSONStorage(() => panelLayoutStorage), }, ), ); diff --git a/packages/ui/src/features/sessions/components/ConversationView.tsx b/packages/ui/src/features/sessions/components/ConversationView.tsx index 275aff3b8c..177d21c40d 100644 --- a/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -37,6 +37,7 @@ import { type VirtualizedListHandle, } from "@posthog/ui/features/sessions/components/VirtualizedList"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; +import { DIFFS_HIGHLIGHTER_OPTIONS } from "@posthog/ui/features/sessions/diffHighlighterOptions"; import { useContextUsage } from "@posthog/ui/features/sessions/hooks/useContextUsage"; import { useConversationItems } from "@posthog/ui/features/sessions/hooks/useConversationItems"; import { useConversationSearch } from "@posthog/ui/features/sessions/hooks/useConversationSearch"; @@ -60,10 +61,6 @@ import { import { Box, Flex, Text } from "@radix-ui/themes"; import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react"; -const DIFFS_HIGHLIGHTER_OPTIONS = { - theme: { dark: "github-dark" as const, light: "github-light" as const }, -}; - export interface ConversationViewProps { events: AcpMessage[]; isPromptPending: boolean | null; diff --git a/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx b/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx index 50a6c9918d..0d1e2f5eca 100644 --- a/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx +++ b/packages/ui/src/features/sessions/components/GeneratingIndicator.tsx @@ -152,7 +152,7 @@ export function GeneratingIndicator({ const startTime = startedAt ?? Date.now(); const interval = setInterval(() => { setElapsed(Math.max(0, Date.now() - startTime - pausedRef.current)); - }, 50); + }, 100); return () => clearInterval(interval); }, [startedAt]); diff --git a/packages/ui/src/features/sessions/components/SessionView.tsx b/packages/ui/src/features/sessions/components/SessionView.tsx index 7f59d8f607..ab832577da 100644 --- a/packages/ui/src/features/sessions/components/SessionView.tsx +++ b/packages/ui/src/features/sessions/components/SessionView.tsx @@ -32,6 +32,7 @@ import { SessionResourcesBar } from "@posthog/ui/features/sessions/components/Se import { SteerQueueToggle } from "@posthog/ui/features/sessions/components/SteerQueueToggle"; import { ThreadView } from "@posthog/ui/features/sessions/components/ThreadView"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; +import { useSessionEventsResidency } from "@posthog/ui/features/sessions/hooks/useSessionEventsResidency"; import { useToggleMessagingMode } from "@posthog/ui/features/sessions/hooks/useToggleMessagingMode"; import { useAdapterForTask, @@ -44,7 +45,7 @@ import { useShowRawLogs, } from "@posthog/ui/features/sessions/sessionViewStore"; import type { Plan } from "@posthog/ui/features/sessions/types"; -import { useSessionForTask } from "@posthog/ui/features/sessions/useSession"; +import { useSessionHandoffInProgress } from "@posthog/ui/features/sessions/useSession"; import { useSettingsStore } from "@posthog/ui/features/settings/settingsStore"; import { useIsWorkspaceCloudRun } from "@posthog/ui/features/workspace/useWorkspace"; import { useConnectivity } from "@posthog/ui/hooks/useConnectivity"; @@ -164,6 +165,7 @@ export function SessionView({ hideInput = false, }: SessionViewProps) { const sessionService = useService(SESSION_SERVICE); + useSessionEventsResidency(taskId); const showRawLogs = useShowRawLogs(); const { setShowRawLogs } = useSessionViewActions(); const pendingTaskPrompt = usePendingTaskPrompt(taskId); @@ -176,8 +178,7 @@ export function SessionView({ const useNewChatThread = useSettingsStore((s) => s.useNewChatThread); const { isOnline } = useConnectivity(); const currentModeId = modeOption?.currentValue; - const handoffInProgress = - useSessionForTask(taskId)?.handoffInProgress ?? false; + const handoffInProgress = useSessionHandoffInProgress(taskId); const showInlineBanner = hasError && errorRetryable && events.length > 0; useEffect(() => { diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index 45fd378ee9..eb5f7f39e3 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -46,10 +46,11 @@ import { import { SessionUpdateView } from "@posthog/ui/features/sessions/components/session-update/SessionUpdateView"; import { UserShellExecuteView } from "@posthog/ui/features/sessions/components/session-update/UserShellExecuteView"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; +import { DIFFS_HIGHLIGHTER_OPTIONS } from "@posthog/ui/features/sessions/diffHighlighterOptions"; import { useConversationItems } from "@posthog/ui/features/sessions/hooks/useConversationItems"; import { useOptimisticItemsForTask, - useSessionForTask, + useSessionIsCloud, } from "@posthog/ui/features/sessions/sessionStore"; import type { UserMessageAttachment } from "@posthog/ui/features/sessions/userMessageTypes"; import { @@ -73,13 +74,8 @@ import { useRef, useState, } from "react"; - import type { ConversationViewProps } from "../ConversationView"; -const DIFFS_HIGHLIGHTER_OPTIONS = { - theme: { dark: "github-dark" as const, light: "github-light" as const }, -}; - /** A row is either a parsed conversation item or a synthesized group of tool calls. */ type ThreadItem = ConversationItem | ToolGroupItem; @@ -570,7 +566,7 @@ export function ChatThread({ ); const optimisticItems = useOptimisticItemsForTask(taskId); - const isCloud = useSessionForTask(taskId)?.isCloud ?? false; + const isCloud = useSessionIsCloud(taskId); const items = useMemo( () => diff --git a/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts b/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts index 84af7898f4..d793141e2c 100644 --- a/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts +++ b/packages/ui/src/features/sessions/components/incrementalConversationItems.test.ts @@ -277,6 +277,77 @@ describe("createIncrementalConversationBuilder", () => { }, ); + // Stream every event (populating the persistent builder), then flip to idle + // so the turn end takes the finalize-in-place path rather than a full rebuild. + it.each(Object.entries(SCENARIOS))( + "finalizes in place equivalently after streaming — %s", + (_name, events) => { + const inc = createIncrementalConversationBuilder(); + for (let k = 1; k <= events.length; k++) { + inc.update(events.slice(0, k), true); + } + expect(normalize(inc.update(events, false))).toEqual( + normalize(buildConversationItems(events, false)), + ); + }, + ); + + it("stays equivalent when streaming resumes after an in-place finalize", () => { + const events = SCENARIOS["multi-turn with tools"]; + const inc = createIncrementalConversationBuilder(); + const firstTurnEnd = 7; // through promptResponseMsg(7, 1) + + for (let k = 1; k <= firstTurnEnd; k++) + inc.update(events.slice(0, k), true); + // Idle after turn 1 → finalize-in-place, which resets the builder. + expect(normalize(inc.update(events.slice(0, firstTurnEnd), false))).toEqual( + normalize(buildConversationItems(events.slice(0, firstTurnEnd), false)), + ); + + // Resume streaming turn 2 on the reset builder, then idle again. + for (let k = firstTurnEnd + 1; k <= events.length; k++) { + inc.update(events.slice(0, k), true); + } + expect(normalize(inc.update(events, false))).toEqual( + normalize(buildConversationItems(events, false)), + ); + }); + + // The idle call arrives with trailing events the builder hasn't seen, so the + // finalize-in-place catch-up loop must process them before finalizing. + it("catches up trailing events that arrive with the idle flip", () => { + const events = SCENARIOS["multi-turn with tools"]; + const inc = createIncrementalConversationBuilder(); + const streamedPrefix = 7; // through promptResponseMsg(7, 1) + + for (let k = 1; k <= streamedPrefix; k++) { + inc.update(events.slice(0, k), true); + } + + expect(normalize(inc.update(events, false))).toEqual( + normalize(buildConversationItems(events, false)), + ); + }); + + // A full rebuild sorts by ts while the incremental builder processed arrival + // order, so out-of-order events must reject finalize-in-place and fall back. + it("falls back to a full rebuild on out-of-order timestamps at idle", () => { + const events = [ + userPromptMsg(1, 1, "hello"), + agentChunk(5, "later "), + agentChunk(3, "earlier "), + promptResponseMsg(6, 1), + ]; + const inc = createIncrementalConversationBuilder(); + for (let k = 1; k <= events.length; k++) { + inc.update(events.slice(0, k), true); + } + + expect(normalize(inc.update(events, false))).toEqual( + normalize(buildConversationItems(events, false)), + ); + }); + it("keeps completed-turn item references stable while the active turn streams", () => { const inc = createIncrementalConversationBuilder(); const base = [ diff --git a/packages/ui/src/features/sessions/components/incrementalConversationItems.ts b/packages/ui/src/features/sessions/components/incrementalConversationItems.ts index 2b6891700f..09b87c330b 100644 --- a/packages/ui/src/features/sessions/components/incrementalConversationItems.ts +++ b/packages/ui/src/features/sessions/components/incrementalConversationItems.ts @@ -5,6 +5,7 @@ import { buildConversationItems, type ConversationItem, createItemBuilder, + finalizeBuilder, type ItemBuilder, markThoughtCompletion, processEvent, @@ -49,9 +50,46 @@ export function createIncrementalConversationBuilder() { ): BuildResult { const debug = options?.showDebugLogs; - // Idle (not streaming): cheap to rebuild, and it sidesteps the speculative - // end-of-stream completions that only `buildConversationItems` resolves. + // Idle (not streaming): finalize the persistent builder in place instead of + // re-parsing every event, but only when the append-only prefix is still + // valid AND events are already in ts-order — a full rebuild sorts, while the + // incremental builder processed in arrival order, so out-of-order events + // must fall back to keep output identical. if (isPromptPending === false) { + let inOrder = true; + for (let i = 1; i < events.length; i++) { + if (events[i].ts < events[i - 1].ts) { + inOrder = false; + break; + } + } + const canFinalizeInPlace = + inOrder && + b !== null && + debug === showDebugLogs && + events.length >= processedCount && + (processedCount === 0 || events[0] === firstEventRef) && + (processedCount === 0 || + events[processedCount - 1] === boundaryEventRef); + + if (canFinalizeInPlace) { + const builder = b as ItemBuilder; + for (let i = processedCount; i < events.length; i++) { + processEvent(builder, events[i], options); + } + finalizeBuilder(builder, isPromptPending); + const result: BuildResult = { + items: builder.items, + lastTurnInfo: readLastTurnInfo(builder), + isCompacting: builder.isCompacting, + completedToolCallCount: builder.completedToolCallCount, + }; + // A finalized builder can't be safely continued; the next streaming + // call rebuilds fresh. + reset(); + return result; + } + reset(); return buildConversationItems(events, isPromptPending, options); } diff --git a/packages/ui/src/features/sessions/components/session-update/CodePreview.tsx b/packages/ui/src/features/sessions/components/session-update/CodePreview.tsx index 3b1d29754f..9859c8ea71 100644 --- a/packages/ui/src/features/sessions/components/session-update/CodePreview.tsx +++ b/packages/ui/src/features/sessions/components/session-update/CodePreview.tsx @@ -1,6 +1,7 @@ import { EditorView } from "@codemirror/view"; import { MultiFileDiff } from "@pierre/diffs/react"; import { compactHomePath, parseImageDataUrl } from "@posthog/shared"; +import { DIFFS_HIGHLIGHTER_OPTIONS } from "@posthog/ui/features/sessions/diffHighlighterOptions"; import { Code } from "@radix-ui/themes"; import { useEffect, useMemo, useRef } from "react"; import { SafeImagePreview } from "../../../../primitives/SafeImagePreview"; @@ -146,10 +147,10 @@ function DiffPreview({ ); const options = useMemo( () => ({ + ...DIFFS_HIGHLIGHTER_OPTIONS, diffStyle: "unified" as const, overflow: "wrap" as const, themeType: (isDarkMode ? "dark" : "light") as "dark" | "light", - theme: { dark: "github-dark" as const, light: "github-light" as const }, disableFileHeader: true, }), [isDarkMode], diff --git a/packages/ui/src/features/sessions/diffHighlighterOptions.ts b/packages/ui/src/features/sessions/diffHighlighterOptions.ts new file mode 100644 index 0000000000..328908e846 --- /dev/null +++ b/packages/ui/src/features/sessions/diffHighlighterOptions.ts @@ -0,0 +1,9 @@ +/** + * Diff highlighter options shared by the session transcript views. + * `tokenizeMaxLineLength` caps tokenization so a minified or single-giant-line + * file can't stall diff highlighting. + */ +export const DIFFS_HIGHLIGHTER_OPTIONS = { + theme: { dark: "github-dark" as const, light: "github-light" as const }, + tokenizeMaxLineLength: 1000, +}; diff --git a/packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts b/packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts index 2c9ec28ef2..fccef985c1 100644 --- a/packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts +++ b/packages/ui/src/features/sessions/hooks/useReturnQueuedMessageToEditor.ts @@ -10,7 +10,7 @@ import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; import { type QueuedMessage, sessionStoreSetters, - useSessionForTask, + useSessionIsCloud, } from "@posthog/ui/features/sessions/sessionStore"; import { useCallback } from "react"; @@ -24,7 +24,7 @@ export function useReturnQueuedMessageToEditor( taskId: string | undefined, ): (message: QueuedMessage) => void { const { requestFocus, setPendingContent } = useDraftStore((s) => s.actions); - const isCloud = useSessionForTask(taskId)?.isCloud ?? false; + const isCloud = useSessionIsCloud(taskId); return useCallback( (message: QueuedMessage) => { diff --git a/packages/ui/src/features/sessions/hooks/useSessionEventsResidency.test.tsx b/packages/ui/src/features/sessions/hooks/useSessionEventsResidency.test.tsx new file mode 100644 index 0000000000..ceb7bf403b --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useSessionEventsResidency.test.tsx @@ -0,0 +1,66 @@ +import { renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const sessionService = vi.hoisted(() => ({ + ensureEventsLoaded: vi.fn().mockResolvedValue(undefined), + scheduleEventEviction: vi.fn(), +})); + +vi.mock("@posthog/di/react", () => ({ + useService: () => sessionService, +})); + +import { useSessionEventsResidency } from "./useSessionEventsResidency"; + +describe("useSessionEventsResidency", () => { + beforeEach(() => { + sessionService.ensureEventsLoaded.mockClear(); + sessionService.scheduleEventEviction.mockClear(); + }); + + it("loads events on mount and schedules eviction on unmount", () => { + const { unmount } = renderHook(() => useSessionEventsResidency("task-1")); + + expect(sessionService.ensureEventsLoaded).toHaveBeenCalledWith("task-1"); + expect(sessionService.scheduleEventEviction).not.toHaveBeenCalled(); + + unmount(); + expect(sessionService.scheduleEventEviction).toHaveBeenCalledWith("task-1"); + }); + + it("does nothing without a taskId", () => { + const { unmount } = renderHook(() => useSessionEventsResidency(undefined)); + unmount(); + + expect(sessionService.ensureEventsLoaded).not.toHaveBeenCalled(); + expect(sessionService.scheduleEventEviction).not.toHaveBeenCalled(); + }); + + it("defers eviction until the last concurrent viewer unmounts", () => { + const first = renderHook(() => useSessionEventsResidency("task-1")); + const second = renderHook(() => useSessionEventsResidency("task-1")); + + first.unmount(); + expect(sessionService.scheduleEventEviction).not.toHaveBeenCalled(); + + second.unmount(); + expect(sessionService.scheduleEventEviction).toHaveBeenCalledTimes(1); + expect(sessionService.scheduleEventEviction).toHaveBeenCalledWith("task-1"); + }); + + it("schedules eviction for the old task when taskId changes", () => { + const { rerender, unmount } = renderHook( + ({ taskId }: { taskId: string }) => useSessionEventsResidency(taskId), + { initialProps: { taskId: "task-1" } }, + ); + + rerender({ taskId: "task-2" }); + expect(sessionService.scheduleEventEviction).toHaveBeenCalledWith("task-1"); + expect(sessionService.ensureEventsLoaded).toHaveBeenLastCalledWith( + "task-2", + ); + + unmount(); + expect(sessionService.scheduleEventEviction).toHaveBeenCalledWith("task-2"); + }); +}); diff --git a/packages/ui/src/features/sessions/hooks/useSessionEventsResidency.ts b/packages/ui/src/features/sessions/hooks/useSessionEventsResidency.ts new file mode 100644 index 0000000000..74878f1b44 --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useSessionEventsResidency.ts @@ -0,0 +1,36 @@ +import { + SESSION_SERVICE, + type SessionService, +} from "@posthog/core/sessions/sessionService"; +import { useService } from "@posthog/di/react"; +import { useEffect } from "react"; + +/** Mounted viewers per taskId, so one view unmounting can't schedule an + * eviction out from under another still-mounted view of the same task. */ +const viewerCounts = new Map(); + +/** + * Ties a task's transcript memory to whether its view is mounted: reloads the + * transcript from disk on view (if it was freed while backgrounded) and + * schedules it to be freed a short while after the last view unmounts. Only + * disconnected background sessions are actually evicted — see + * {@link SessionService.scheduleEventEviction}. + */ +export function useSessionEventsResidency(taskId: string | undefined): void { + const sessionService = useService(SESSION_SERVICE); + + useEffect(() => { + if (!taskId) return; + viewerCounts.set(taskId, (viewerCounts.get(taskId) ?? 0) + 1); + void sessionService.ensureEventsLoaded(taskId); + return () => { + const remaining = (viewerCounts.get(taskId) ?? 1) - 1; + if (remaining > 0) { + viewerCounts.set(taskId, remaining); + return; + } + viewerCounts.delete(taskId); + sessionService.scheduleEventEviction(taskId); + }; + }, [taskId, sessionService]); +} diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index cb437bc724..fd9f75b653 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -4767,6 +4767,8 @@ describe("SessionService", () => { }, }; onData(echo); + // Streamed events are buffered and flushed on a frame timer; let it run. + await new Promise((resolve) => setTimeout(resolve, 25)); if (steer) { expect(mockSessionStoreSetters.appendEvents).toHaveBeenCalledWith( diff --git a/packages/ui/src/features/sessions/sessionStore.ts b/packages/ui/src/features/sessions/sessionStore.ts index b9432901f0..d266829b6f 100644 --- a/packages/ui/src/features/sessions/sessionStore.ts +++ b/packages/ui/src/features/sessions/sessionStore.ts @@ -86,6 +86,8 @@ export { usePendingPermissionsForTask, useQueuedMessagesForTask, useSessionForTask, + useSessionHandoffInProgress, + useSessionIsCloud, useSessions, useThoughtLevelConfigOptionForTask, } from "./useSession"; diff --git a/packages/ui/src/features/sessions/useSession.ts b/packages/ui/src/features/sessions/useSession.ts index 0352ee1095..00a9ccfaed 100644 --- a/packages/ui/src/features/sessions/useSession.ts +++ b/packages/ui/src/features/sessions/useSession.ts @@ -150,3 +150,30 @@ export const useAdapterForTask = ( return s.sessions[taskRunId]?.adapter; }); }; + +/** + * Whether a task's session is a cloud run. A primitive selector, so consumers + * that only need this flag don't re-render on every streamed event the way + * reading the whole session via {@link useSessionForTask} would. + */ +export const useSessionIsCloud = (taskId: string | undefined): boolean => { + return useSessionStore((s) => { + if (!taskId) return false; + const taskRunId = s.taskIdIndex[taskId]; + if (!taskRunId) return false; + return s.sessions[taskRunId]?.isCloud ?? false; + }); +}; + +/** Whether a cloud handoff is in progress for a task. Primitive selector — see + * {@link useSessionIsCloud}. */ +export const useSessionHandoffInProgress = ( + taskId: string | undefined, +): boolean => { + return useSessionStore((s) => { + if (!taskId) return false; + const taskRunId = s.taskIdIndex[taskId]; + if (!taskRunId) return false; + return s.sessions[taskRunId]?.handoffInProgress ?? false; + }); +}; diff --git a/packages/ui/src/features/sidebar/useSidebarData.ts b/packages/ui/src/features/sidebar/useSidebarData.ts index a040204883..8aac61fb8b 100644 --- a/packages/ui/src/features/sidebar/useSidebarData.ts +++ b/packages/ui/src/features/sidebar/useSidebarData.ts @@ -18,12 +18,12 @@ import type { AppView } from "@posthog/ui/router/useAppView"; import { useEffect, useMemo, useRef } from "react"; import { useArchivedTaskIds } from "../archive/useArchivedTaskIds"; import { useProvisioningStore } from "../provisioning/store"; -import { useSessions } from "../sessions/sessionStore"; import { useSuspendedTaskIds } from "../suspension/useSuspendedTaskIds"; import { useSlackTasks, useTaskSummaries, useTasks } from "../tasks/useTasks"; import { useWorkspaces } from "../workspace/useWorkspace"; import { useSidebarStore } from "./sidebarStore"; import { usePinnedTasks } from "./usePinnedTasks"; +import { useSidebarSessionMap } from "./useSidebarSessionMap"; import { useTaskViewed } from "./useTaskViewed"; export type { SidebarData, TaskData, TaskGroup }; @@ -41,7 +41,7 @@ export function useSidebarData({ const archivedTaskIds = useArchivedTaskIds(); const suspendedTaskIds = useSuspendedTaskIds(); const provisioningTaskIds = useProvisioningStore((s) => s.activeTasks); - const sessions = useSessions(); + const sessionByTaskId = useSidebarSessionMap(); const { timestamps } = useTaskViewed(); const historyVisibleCount = useSidebarStore( (state) => state.historyVisibleCount, @@ -140,16 +140,6 @@ export function useSidebarData({ const activeTaskId = activeView.type === "task-detail" ? (activeView.taskId ?? null) : null; - const sessionByTaskId = useMemo(() => { - const map = new Map(); - for (const session of Object.values(sessions)) { - if (session.taskId) { - map.set(session.taskId, session); - } - } - return map; - }, [sessions]); - const taskData = useMemo( () => allTasks.map((task) => diff --git a/packages/ui/src/features/sidebar/useSidebarSessionMap.test.tsx b/packages/ui/src/features/sidebar/useSidebarSessionMap.test.tsx new file mode 100644 index 0000000000..d4dbaa0249 --- /dev/null +++ b/packages/ui/src/features/sidebar/useSidebarSessionMap.test.tsx @@ -0,0 +1,60 @@ +import { sessionStoreSetters } from "@posthog/core/sessions/sessionStore"; +import type { AcpMessage, AgentSession } from "@posthog/shared"; +import { act, renderHook } from "@testing-library/react"; +import { afterEach, describe, expect, it } from "vitest"; +import { useSidebarSessionMap } from "./useSidebarSessionMap"; + +const RUN_ID = "run-1"; +const TASK_ID = "task-1"; + +function seedSession() { + sessionStoreSetters.setSession({ + taskRunId: RUN_ID, + taskId: TASK_ID, + events: [], + messageQueue: [], + pendingPermissions: new Map(), + isPromptPending: false, + } as unknown as AgentSession); +} + +afterEach(() => { + sessionStoreSetters.removeSession(RUN_ID); +}); + +describe("useSidebarSessionMap", () => { + it("does not re-render when only events are appended", () => { + seedSession(); + let renders = 0; + renderHook(() => { + renders++; + return useSidebarSessionMap(); + }); + const baseline = renders; + + act(() => { + sessionStoreSetters.appendEvents(RUN_ID, [ + { ts: 1, message: {} } as unknown as AcpMessage, + ]); + }); + + expect(renders).toBe(baseline); + }); + + it("re-renders when a sidebar-relevant field changes", () => { + seedSession(); + let renders = 0; + const { result } = renderHook(() => { + renders++; + return useSidebarSessionMap(); + }); + const baseline = renders; + + act(() => { + sessionStoreSetters.updateSession(RUN_ID, { isPromptPending: true }); + }); + + expect(renders).toBeGreaterThan(baseline); + expect(result.current.get(TASK_ID)?.isPromptPending).toBe(true); + }); +}); diff --git a/packages/ui/src/features/sidebar/useSidebarSessionMap.ts b/packages/ui/src/features/sidebar/useSidebarSessionMap.ts new file mode 100644 index 0000000000..60018438ed --- /dev/null +++ b/packages/ui/src/features/sidebar/useSidebarSessionMap.ts @@ -0,0 +1,27 @@ +import { computeSidebarSessionSignature } from "@posthog/core/sidebar/buildSidebarData"; +import type { AgentSession } from "@posthog/shared"; +import { useMemo } from "react"; +import { useSessionStore } from "../sessions/sessionStore"; + +/** + * `taskId → session` map for the sidebar, rebuilt only when a sidebar-relevant + * session field changes — not on every streamed event. The equality function + * compares just the fields {@link computeSidebarSessionSignature} covers, so the + * subscription (and the root-mounted sidebar) ignores the appends that fire on + * every token during a turn. + */ +export function useSidebarSessionMap(): Map { + const sessions = useSessionStore( + (s) => s.sessions, + (a, b) => + computeSidebarSessionSignature(a) === computeSidebarSessionSignature(b), + ); + + return useMemo(() => { + const map = new Map(); + for (const session of Object.values(sessions)) { + if (session.taskId) map.set(session.taskId, session); + } + return map; + }, [sessions]); +} diff --git a/packages/ui/src/features/task-detail/components/TabContentRenderer.tsx b/packages/ui/src/features/task-detail/components/TabContentRenderer.tsx index ed1e2a4d80..72d23465d6 100644 --- a/packages/ui/src/features/task-detail/components/TabContentRenderer.tsx +++ b/packages/ui/src/features/task-detail/components/TabContentRenderer.tsx @@ -1,7 +1,9 @@ import type { Task } from "@posthog/shared/domain-types"; import { CodeEditorPanel } from "../../code-editor/components/CodeEditorPanel"; -import { CloudReviewPage } from "../../code-review/components/CloudReviewPage"; -import { ReviewPage } from "../../code-review/components/ReviewPage"; +import { + LazyCloudReviewPage as CloudReviewPage, + LazyReviewPage as ReviewPage, +} from "../../code-review/components/LazyReviewPages"; import type { Tab } from "../../panels/panelTypes"; import { useIsWorkspaceCloudRun } from "../../workspace/useWorkspace"; import { ActionPanel } from "./ActionPanel"; diff --git a/packages/ui/src/features/task-detail/components/TaskDetail.tsx b/packages/ui/src/features/task-detail/components/TaskDetail.tsx index 1486fcceb9..5d9e89fa4a 100644 --- a/packages/ui/src/features/task-detail/components/TaskDetail.tsx +++ b/packages/ui/src/features/task-detail/components/TaskDetail.tsx @@ -6,8 +6,10 @@ import { useBlurOnEscape } from "../../../hooks/useBlurOnEscape"; import { useSetHeaderContent } from "../../../hooks/useSetHeaderContent"; import { logger } from "../../../shell/logger"; import { ChannelBreadcrumb } from "../../canvas/components/ChannelBreadcrumb"; -import { CloudReviewPage } from "../../code-review/components/CloudReviewPage"; -import { ReviewPage } from "../../code-review/components/ReviewPage"; +import { + LazyCloudReviewPage as CloudReviewPage, + LazyReviewPage as ReviewPage, +} from "../../code-review/components/LazyReviewPages"; import { useReviewNavigationStore } from "../../code-review/reviewNavigationStore"; import { FilePicker } from "../../command/FilePicker"; import { useRepoFileWatcher } from "../../file-watcher/useRepoFileWatcher"; diff --git a/packages/ui/src/utils/syntax-highlight.test.ts b/packages/ui/src/utils/syntax-highlight.test.ts new file mode 100644 index 0000000000..9e2f963c5c --- /dev/null +++ b/packages/ui/src/utils/syntax-highlight.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "vitest"; +import { highlightSyntax } from "./syntax-highlight"; + +describe("highlightSyntax", () => { + it("returns segments whose text reconstructs the original code", () => { + const code = "const x = 1;\nconst y = 2;"; + const segments = highlightSyntax(code, "typescript", true); + expect(segments).not.toBeNull(); + expect(segments?.map((s) => s.text).join("")).toBe(code); + }); + + it("returns the cached array on a repeated identical call", () => { + const code = "def add(a, b):\n return a + b"; + const first = highlightSyntax(code, "python", true); + const second = highlightSyntax(code, "python", true); + expect(first).not.toBeNull(); + expect(second).toBe(first); + }); + + it("caches per theme — light and dark are distinct results", () => { + const code = "let z = 3;"; + const dark = highlightSyntax(code, "javascript", true); + const light = highlightSyntax(code, "javascript", false); + expect(dark).not.toBe(light); + }); + + it("returns null for an unsupported language", () => { + expect(highlightSyntax("whatever", "brainfuck", true)).toBeNull(); + }); + + it("evicts the oldest entry once the cache is full", () => { + const code = "const evicted = true;"; + const first = highlightSyntax(code, "javascript", true); + + // 256 distinct inserts push everything older out of the bounded cache. + for (let i = 0; i < 256; i++) { + highlightSyntax(`const filler${i} = ${i};`, "javascript", true); + } + + const again = highlightSyntax(code, "javascript", true); + expect(again).not.toBe(first); + expect(again).toEqual(first); + }); +}); diff --git a/packages/ui/src/utils/syntax-highlight.ts b/packages/ui/src/utils/syntax-highlight.ts index d0b3c35843..79fd2cef61 100644 --- a/packages/ui/src/utils/syntax-highlight.ts +++ b/packages/ui/src/utils/syntax-highlight.ts @@ -151,6 +151,28 @@ export interface HighlightSegment { color?: string; } +/** + * Parsed output cache keyed by (theme, language, length, content hash). The + * per-component useMemo only survives that instance, so virtualized scroll + * re-parses a code block every time it remounts. This bounded LRU makes + * remounts free. The code is stored alongside the segments so a hash collision + * is detected on lookup rather than returning another snippet's output. + */ +const MAX_HIGHLIGHT_CACHE_ENTRIES = 256; +const highlightCache = new Map< + string, + { code: string; segments: HighlightSegment[] } +>(); + +function hashCode(text: string): number { + let hash = 0x811c9dc5; + for (let i = 0; i < text.length; i++) { + hash ^= text.charCodeAt(i); + hash = Math.imul(hash, 0x01000193); + } + return hash >>> 0; +} + export function highlightSyntax( code: string, language: string, @@ -159,6 +181,14 @@ export function highlightSyntax( const parser = getParser(language); if (!parser) return null; + const cacheKey = `${isDark ? "d" : "l"}:${language}:${code.length}:${hashCode(code).toString(36)}`; + const cached = highlightCache.get(cacheKey); + if (cached && cached.code === code) { + highlightCache.delete(cacheKey); + highlightCache.set(cacheKey, cached); + return cached.segments; + } + const tree = parser.parse(code); const palette = isDark ? darkPalette : lightPalette; const segments: HighlightSegment[] = []; @@ -177,5 +207,11 @@ export function highlightSyntax( }, ); + highlightCache.set(cacheKey, { code, segments }); + if (highlightCache.size > MAX_HIGHLIGHT_CACHE_ENTRIES) { + const oldest = highlightCache.keys().next().value; + if (oldest !== undefined) highlightCache.delete(oldest); + } + return segments; } diff --git a/packages/workspace-server/src/services/local-logs/identifiers.ts b/packages/workspace-server/src/services/local-logs/identifiers.ts index 9b978898b3..212ef8155b 100644 --- a/packages/workspace-server/src/services/local-logs/identifiers.ts +++ b/packages/workspace-server/src/services/local-logs/identifiers.ts @@ -3,5 +3,14 @@ export const LOGS_SERVICE = Symbol.for("posthog.workspace.logsService"); export interface ILogsService { fetchS3Logs(logUrl: string): Promise; readLocalLogs(taskRunId: string): Promise; + /** + * Read only the last `maxBytes` of the log for a fast initial paint. Returns + * `truncated: true` when older history was skipped (the partial first line is + * dropped). `null` if there's no local log. + */ + readLocalLogsTail( + taskRunId: string, + maxBytes: number, + ): Promise<{ content: string; truncated: boolean } | null>; writeLocalLogs(taskRunId: string, content: string): Promise; } diff --git a/packages/workspace-server/src/services/local-logs/schemas.ts b/packages/workspace-server/src/services/local-logs/schemas.ts index 15f59fc819..708671d973 100644 --- a/packages/workspace-server/src/services/local-logs/schemas.ts +++ b/packages/workspace-server/src/services/local-logs/schemas.ts @@ -6,6 +6,14 @@ export const fetchS3LogsOutput = z.string().nullable(); export const readLocalLogsInput = z.object({ taskRunId: z.string().min(1) }); export const readLocalLogsOutput = z.string().nullable(); +export const readLocalLogsTailInput = z.object({ + taskRunId: z.string().min(1), + maxBytes: z.number().int().positive(), +}); +export const readLocalLogsTailOutput = z + .object({ content: z.string(), truncated: z.boolean() }) + .nullable(); + export const writeLocalLogsInput = z.object({ taskRunId: z.string().min(1), content: z.string(), diff --git a/packages/workspace-server/src/services/local-logs/service.ts b/packages/workspace-server/src/services/local-logs/service.ts index 0ded299537..058e142785 100644 --- a/packages/workspace-server/src/services/local-logs/service.ts +++ b/packages/workspace-server/src/services/local-logs/service.ts @@ -52,6 +52,43 @@ export class LocalLogsService implements ILogsService { } } + async readLocalLogsTail( + taskRunId: string, + maxBytes: number, + ): Promise<{ content: string; truncated: boolean } | null> { + const logPath = this.getLocalLogPath(taskRunId); + try { + const stat = await fs.promises.stat(logPath); + if (stat.size <= maxBytes) { + return { + content: await fs.promises.readFile(logPath, "utf-8"), + truncated: false, + }; + } + const handle = await fs.promises.open(logPath, "r"); + try { + // Read one extra byte before the window: a newline there means the + // window already starts on a whole line. Otherwise the first line is + // a fragment (and may start with a broken multi-byte char) — drop + // everything up to the first newline so only whole ndjson lines + // remain. + const start = stat.size - maxBytes - 1; + const buf = Buffer.alloc(maxBytes + 1); + const { bytesRead } = await handle.read(buf, 0, maxBytes + 1, start); + const raw = buf.toString("utf-8", 1, bytesRead); + if (buf[0] === 0x0a) { + return { content: raw, truncated: true }; + } + const nl = raw.indexOf("\n"); + return { content: nl >= 0 ? raw.slice(nl + 1) : "", truncated: true }; + } finally { + await handle.close(); + } + } catch { + return null; + } + } + writeLocalLogs(taskRunId: string, content: string): Promise { const existing = this.writes.get(taskRunId); if (existing) { diff --git a/packages/workspace-server/src/services/local-logs/serviceTail.test.ts b/packages/workspace-server/src/services/local-logs/serviceTail.test.ts new file mode 100644 index 0000000000..446d96ed92 --- /dev/null +++ b/packages/workspace-server/src/services/local-logs/serviceTail.test.ts @@ -0,0 +1,79 @@ +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { LocalLogsService } from "./service"; + +const RUN = "run-tail"; + +describe("LocalLogsService.readLocalLogsTail", () => { + let tmpHome: string; + + beforeEach(() => { + tmpHome = fs.mkdtempSync(path.join(os.tmpdir(), "phlogs-")); + vi.spyOn(os, "homedir").mockReturnValue(tmpHome); + fs.mkdirSync(path.join(tmpHome, ".posthog-code", "sessions", RUN), { + recursive: true, + }); + }); + + afterEach(() => { + vi.restoreAllMocks(); + fs.rmSync(tmpHome, { recursive: true, force: true }); + }); + + const logPath = () => + path.join(tmpHome, ".posthog-code", "sessions", RUN, "logs.ndjson"); + + it("returns the whole file untruncated when it's under maxBytes", async () => { + const content = "line1\nline2\nline3\n"; + fs.writeFileSync(logPath(), content); + + const res = await new LocalLogsService().readLocalLogsTail(RUN, 1_000_000); + + expect(res).toEqual({ content, truncated: false }); + }); + + it("returns only the tail, dropping the partial first line, when over maxBytes", async () => { + const lines = Array.from( + { length: 1000 }, + (_, i) => `{"i":${i},"pad":"${"x".repeat(200)}"}`, + ); + fs.writeFileSync(logPath(), `${lines.join("\n")}\n`); + + const res = await new LocalLogsService().readLocalLogsTail(RUN, 5000); + + expect(res?.truncated).toBe(true); + const tailLines = res?.content.trim().split("\n") ?? []; + // Every retained line is a whole, parseable ndjson entry (no fragment). + for (const line of tailLines) { + expect(() => JSON.parse(line)).not.toThrow(); + } + // It's the suffix of the file — ends with the last written line. + expect(tailLines.at(-1)).toBe(lines.at(-1)); + // It's a strict tail, not the whole file. + expect(tailLines.length).toBeLessThan(lines.length); + }); + + it("keeps the whole first line when the window starts on a line boundary", async () => { + fs.writeFileSync(logPath(), "aaaa\nbbbb\ncccc\n"); + + const res = await new LocalLogsService().readLocalLogsTail(RUN, 10); + + expect(res).toEqual({ content: "bbbb\ncccc\n", truncated: true }); + }); + + it("returns empty content when a single line exceeds maxBytes", async () => { + fs.writeFileSync(logPath(), `{"pad":"${"x".repeat(500)}"}\n`); + + const res = await new LocalLogsService().readLocalLogsTail(RUN, 100); + + expect(res).toEqual({ content: "", truncated: true }); + }); + + it("returns null when the log doesn't exist", async () => { + expect( + await new LocalLogsService().readLocalLogsTail("missing", 1000), + ).toBeNull(); + }); +}); diff --git a/packages/workspace-server/src/trpc.ts b/packages/workspace-server/src/trpc.ts index 7885d44c27..2769c52014 100644 --- a/packages/workspace-server/src/trpc.ts +++ b/packages/workspace-server/src/trpc.ts @@ -132,6 +132,8 @@ import { deleteLocalLogCacheInput, readLocalLogsInput, readLocalLogsOutput, + readLocalLogsTailInput, + readLocalLogsTailOutput, seedLocalLogsInput, writeLocalLogsInput, } from "./services/local-logs/schemas"; @@ -854,6 +856,13 @@ export function createAppRouter({ localLogsService().readLocalLogs(input.taskRunId), ), + readTail: t.procedure + .input(readLocalLogsTailInput) + .output(readLocalLogsTailOutput) + .query(({ input }) => + localLogsService().readLocalLogsTail(input.taskRunId, input.maxBytes), + ), + write: t.procedure .input(writeLocalLogsInput) .mutation(({ input }) =>