diff --git a/packages/ui/src/components/session/context-usage-panel.tsx b/packages/ui/src/components/session/context-usage-panel.tsx index 3c515805..7c69ea1f 100644 --- a/packages/ui/src/components/session/context-usage-panel.tsx +++ b/packages/ui/src/components/session/context-usage-panel.tsx @@ -1,13 +1,9 @@ -import { createEffect, createMemo, Show, untrack, type Component } from "solid-js" -import { getChildSessions, getSessionFamily, getSessionInfo, loadMessages, sessionInfoByInstance } from "../../stores/sessions" +import { createMemo, Show, type Component } from "solid-js" +import { getChildSessions, getSessionInfo, getThreadTotals } from "../../stores/sessions" import { formatTokenTotal } from "../../lib/formatters" -import { computeThreadTotals } from "../../lib/thread-totals" import { useI18n } from "../../lib/i18n" -import { getLogger } from "../../lib/logger" import { useConfig } from "../../stores/preferences" -const log = getLogger("session") - interface ContextUsagePanelProps { instanceId: string sessionId: string @@ -39,32 +35,9 @@ const ContextUsagePanel: Component = (props) => { const children = createMemo(() => getChildSessions(props.instanceId, props.sessionId)) const hasChildren = createMemo(() => children().length > 0) - const instanceInfoMap = createMemo(() => - sessionInfoByInstance().get(props.instanceId), - ) - - createEffect(() => { - if (!showUsage()) return - const instanceId = props.instanceId - const childSessions = children() - untrack(() => { - for (const child of childSessions) { - loadMessages(instanceId, child.id, { skipDiff: true }).catch((error) => - log.error("Failed to load child session messages", { - instanceId, - sessionId: child.id, - error, - }), - ) - } - }) - }) - const threadTotals = createMemo(() => { if (!hasChildren()) return { cost: 0, inputTokens: 0, outputTokens: 0, reasoningTokens: 0 } - const map = instanceInfoMap() - const family = getSessionFamily(props.instanceId, props.sessionId) - return computeThreadTotals(family, map) + return getThreadTotals(props.instanceId, props.sessionId) ?? { cost: 0, inputTokens: 0, outputTokens: 0, reasoningTokens: 0 } }) const totalCostDisplay = createMemo(() => `$${threadTotals().cost.toFixed(2)}`) diff --git a/packages/ui/src/stores/message-v2/session-info.ts b/packages/ui/src/stores/message-v2/session-info.ts index 50bbf697..fc157510 100644 --- a/packages/ui/src/stores/message-v2/session-info.ts +++ b/packages/ui/src/stores/message-v2/session-info.ts @@ -1,6 +1,6 @@ import type { Provider } from "../../types/session" import { DEFAULT_MODEL_OUTPUT_LIMIT } from "../session-models" -import { providers, sessions, sessionInfoByInstance, setSessionInfoByInstance } from "../session-state" +import { providers, sessions, sessionInfoByInstance, setSessionInfoByInstance, updateThreadTotalsForSession } from "../session-state" import { messageStoreBus } from "./bus" import type { SessionUsageState } from "./types" @@ -140,4 +140,6 @@ export function updateSessionInfo(instanceId: string, sessionId: string): void { next.set(instanceId, instanceInfo) return next }) + + updateThreadTotalsForSession(instanceId, sessionId) } diff --git a/packages/ui/src/stores/session-api.ts b/packages/ui/src/stores/session-api.ts index dad131e4..bf88185e 100644 --- a/packages/ui/src/stores/session-api.ts +++ b/packages/ui/src/stores/session-api.ts @@ -25,6 +25,7 @@ import { setLoading, cleanupBlankSessions, syncInstanceSessionIndicator, + updateThreadTotalsForParent, } from "./session-state" import { DEFAULT_MODEL_OUTPUT_LIMIT, getDefaultModel, isModelValid } from "./session-models" import { normalizeMessagePart } from "./message-v2/normalizers" @@ -45,6 +46,7 @@ import { const log = getLogger("api") const pendingSessionDiffFetches = new Map>() +const pendingSessionChildrenFetches = new Map>() async function loadSessionDiff(instanceId: string, sessionId: string, force = false): Promise { if (!instanceId || !sessionId) return @@ -218,6 +220,12 @@ async function fetchSessions(instanceId: string): Promise { pruneDraftPrompts(instanceId, new Set(sessionMap.keys())) + + const parentIds = Array.from(sessionMap.values()) + .filter((session) => session.parentId === null) + .map((session) => session.id) + + await Promise.all(parentIds.map((parentId) => fetchSessionChildren(instanceId, parentId))) } catch (error) { log.error("Failed to fetch sessions:", error) throw error @@ -230,6 +238,96 @@ async function fetchSessions(instanceId: string): Promise { } } +function toClientSession(instanceId: string, apiSession: any, existingSession?: Session): Session { + return { + id: apiSession.id, + instanceId, + title: apiSession.title || existingSession?.title || "Untitled", + parentId: apiSession.parentID || null, + agent: existingSession?.agent ?? "", + model: existingSession?.model ?? { providerId: "", modelId: "" }, + status: existingSession?.status ?? "idle", + retry: existingSession?.retry ?? null, + idleSince: existingSession?.idleSince ?? null, + version: apiSession.version, + time: { + ...apiSession.time, + }, + revert: apiSession.revert + ? { + messageID: apiSession.revert.messageID, + partID: apiSession.revert.partID, + snapshot: apiSession.revert.snapshot, + diff: apiSession.revert.diff, + } + : existingSession?.revert, + diff: existingSession?.diff, + pendingPermission: existingSession?.pendingPermission, + pendingQuestion: existingSession?.pendingQuestion, + } +} + +async function fetchSessionChildren(instanceId: string, parentSessionId: string): Promise { + if (!instanceId || !parentSessionId) return [] + + const key = `${instanceId}:${parentSessionId}` + const pending = pendingSessionChildrenFetches.get(key) + if (pending) return pending + + const promise = (async () => { + const instance = instances().get(instanceId) + if (!instance || !instance.client) { + throw new Error("Instance not ready") + } + + const worktreeSlug = getWorktreeSlugForSession(instanceId, parentSessionId) + const client = getOrCreateWorktreeClient(instanceId, worktreeSlug) + + log.info(`[HTTP] GET /session/{sessionID}/children for instance ${instanceId}`, { sessionId: parentSessionId }) + const apiChildren = await requestData( + client.session.children({ sessionID: parentSessionId }), + "session.children", + ) + + if (!Array.isArray(apiChildren)) return [] + + const currentSessions = sessions().get(instanceId) + const children = apiChildren.map((apiSession) => toClientSession(instanceId, apiSession, currentSessions?.get(apiSession.id))) + + setSessions((prev) => { + const next = new Map(prev) + const instanceSessions = new Map(next.get(instanceId)) + for (const child of children) { + instanceSessions.set(child.id, child) + } + next.set(instanceId, instanceSessions) + return next + }) + + syncInstanceSessionIndicator(instanceId) + updateThreadTotalsForParent(instanceId, parentSessionId) + + if (messagesLoaded().get(instanceId)?.has(parentSessionId)) { + for (const child of children) { + void loadMessages(instanceId, child.id, { skipDiff: true, skipChildren: true }).catch((error) => + log.error("Failed to load child session messages", { + instanceId, + sessionId: child.id, + parentSessionId, + error, + }), + ) + } + } + + return children + })() + + pendingSessionChildrenFetches.set(key, promise) + void promise.finally(() => pendingSessionChildrenFetches.delete(key)) + return promise +} + async function createSession(instanceId: string, agent?: string): Promise { const instance = instances().get(instanceId) if (!instance || !instance.client) { @@ -598,10 +696,11 @@ async function fetchProviders(instanceId: string): Promise { async function loadMessages( instanceId: string, sessionId: string, - options?: { force?: boolean; skipDiff?: boolean }, + options?: { force?: boolean; skipDiff?: boolean; skipChildren?: boolean }, ): Promise { const force = options?.force ?? false const skipDiff = options?.skipDiff ?? false + const skipChildren = options?.skipChildren ?? false if (force) { setMessagesLoaded((prev) => { @@ -783,6 +882,19 @@ async function loadMessages( } updateSessionInfo(instanceId, sessionId) + + if (!skipChildren && session.parentId === null) { + for (const child of getChildSessions(instanceId, sessionId)) { + void loadMessages(instanceId, child.id, { skipDiff: true, skipChildren: true }).catch((error) => + log.error("Failed to load child session messages", { + instanceId, + sessionId: child.id, + parentSessionId: sessionId, + error, + }), + ) + } + } } export { @@ -792,6 +904,7 @@ export { fetchProviders, fetchSessions, + fetchSessionChildren, forkSession, loadMessages, } diff --git a/packages/ui/src/stores/session-state.ts b/packages/ui/src/stores/session-state.ts index f5e82cbd..726f7cb8 100644 --- a/packages/ui/src/stores/session-state.ts +++ b/packages/ui/src/stores/session-state.ts @@ -10,6 +10,7 @@ import { getLogger } from "../lib/logger" import { requestData } from "../lib/opencode-api" import { getOrCreateWorktreeClient, getWorktreeSlugForSession } from "./worktrees" import { tGlobal } from "../lib/i18n" +import { computeThreadTotals, type ThreadTotals } from "../lib/thread-totals" const log = getLogger("session") @@ -47,6 +48,7 @@ const [loading, setLoading] = createSignal({ const [messagesLoaded, setMessagesLoaded] = createSignal>>(new Map()) const [sessionInfoByInstance, setSessionInfoByInstance] = createSignal>>(new Map()) +const [threadTotalsByInstance, setThreadTotalsByInstance] = createSignal>>(new Map()) const [expandedSessionParents, setExpandedSessionParents] = createSignal>>(new Map()) @@ -648,6 +650,29 @@ function getSessionInfo(instanceId: string, sessionId: string): SessionInfo | un return sessionInfoByInstance().get(instanceId)?.get(sessionId) } +function getThreadTotals(instanceId: string, parentSessionId: string): ThreadTotals | undefined { + return threadTotalsByInstance().get(instanceId)?.get(parentSessionId) +} + +function updateThreadTotalsForParent(instanceId: string, parentSessionId: string): void { + const family = getSessionFamily(instanceId, parentSessionId) + const totals = computeThreadTotals(family, sessionInfoByInstance().get(instanceId)) + + setThreadTotalsByInstance((prev) => { + const next = new Map(prev) + const instanceTotals = new Map(next.get(instanceId)) + instanceTotals.set(parentSessionId, totals) + next.set(instanceId, instanceTotals) + return next + }) +} + +function updateThreadTotalsForSession(instanceId: string, sessionId: string): void { + const session = sessions().get(instanceId)?.get(sessionId) + if (!session) return + updateThreadTotalsForParent(instanceId, session.parentId ?? session.id) +} + async function isBlankSession(session: Session, instanceId: string, fetchIfNeeded = false): Promise { const created = session.time?.created || 0 const updated = session.time?.updated || 0 @@ -774,6 +799,10 @@ export { setMessagesLoaded, sessionInfoByInstance, setSessionInfoByInstance, + threadTotalsByInstance, + getThreadTotals, + updateThreadTotalsForParent, + updateThreadTotalsForSession, getSessionDraftPrompt, setSessionDraftPrompt, clearSessionDraftPrompt, diff --git a/packages/ui/src/stores/sessions.ts b/packages/ui/src/stores/sessions.ts index 96ac43e6..b64b04a7 100644 --- a/packages/ui/src/stores/sessions.ts +++ b/packages/ui/src/stores/sessions.ts @@ -18,6 +18,7 @@ import { getSessionFamily, getSessionInfo, getSessionThreads, + getThreadTotals, getSessions, getVisibleSessionIds, isSessionBusy, @@ -45,6 +46,7 @@ import { fetchAgents, fetchProviders, fetchSessions, + fetchSessionChildren, forkSession, loadMessages, } from "./session-api" @@ -109,6 +111,7 @@ export { fetchAgents, fetchProviders, fetchSessions, + fetchSessionChildren, forkSession, getActiveParentSession, getActiveSession, @@ -119,6 +122,7 @@ export { getSessionFamily, getSessionInfo, getSessionThreads, + getThreadTotals, getSessions, getVisibleSessionIds, isSessionBusy,