fix session performance issue from large diffs

This commit is contained in:
Dax Raad
2025-11-04 13:32:56 -05:00
parent f9af9fc221
commit 7a7060ef15
6 changed files with 101 additions and 17 deletions

View File

@@ -758,6 +758,34 @@ export namespace Server {
return c.json(messages)
},
)
.get(
"/session/:id/diff",
describeRoute({
description: "Get the diff for this session",
operationId: "session.diff",
responses: {
200: {
description: "List of diffs",
content: {
"application/json": {
schema: resolver(Snapshot.FileDiff.array()),
},
},
},
...errors(400, 404),
},
}),
validator(
"param",
z.object({
id: z.string().meta({ description: "Session ID" }),
}),
),
async (c) => {
const diff = await Session.diff(c.req.valid("param").id)
return c.json(diff)
},
)
.get(
"/session/:id/message/:messageID",
describeRoute({

View File

@@ -15,8 +15,8 @@ import { MessageV2 } from "./message-v2"
import { Instance } from "../project/instance"
import { SessionPrompt } from "./prompt"
import { fn } from "@/util/fn"
import { Snapshot } from "@/snapshot"
import { Command } from "../command"
import { Snapshot } from "@/snapshot"
export namespace Session {
const log = Log.create({ service: "session" })
@@ -42,7 +42,9 @@ export namespace Session {
parentID: Identifier.schema("session").optional(),
summary: z
.object({
diffs: Snapshot.FileDiff.array(),
additions: z.number(),
deletions: z.number(),
diffs: Snapshot.FileDiff.array().optional(),
})
.optional(),
share: z
@@ -258,6 +260,11 @@ export namespace Session {
return result
}
export const diff = fn(Identifier.schema("session"), async (sessionID) => {
const diffs = await Storage.read<Snapshot.FileDiff[]>(["session_diff", sessionID])
return diffs ?? []
})
export const messages = fn(Identifier.schema("session"), async (sessionID) => {
const result = [] as MessageV2.WithParts[]
for (const p of await Storage.list(["message", sessionID])) {

View File

@@ -11,6 +11,7 @@ import { SystemPrompt } from "./system"
import { Log } from "@/util/log"
import path from "path"
import { Instance } from "@/project/instance"
import { Storage } from "@/storage/storage"
export namespace SessionSummary {
const log = Log.create({ service: "session.summary" })
@@ -44,9 +45,11 @@ export namespace SessionSummary {
)
await Session.update(input.sessionID, (draft) => {
draft.summary = {
diffs,
additions: diffs.reduce((sum, x) => sum + x.additions, 0),
deletions: diffs.reduce((sum, x) => sum + x.deletions, 0),
}
})
await Storage.write(["session_diff", input.sessionID], diffs)
}
async function summarizeMessage(input: { messageID: string; messages: MessageV2.WithParts[] }) {

View File

@@ -85,7 +85,9 @@ export namespace Storage {
const session = await Bun.file(sessionFile).json()
await Bun.write(dest, JSON.stringify(session))
log.info(`migrating messages for session ${session.id}`)
for await (const msgFile of new Bun.Glob(`storage/session/message/${session.id}/*.json`).scan({
for await (const msgFile of new Bun.Glob(
`storage/session/message/${session.id}/*.json`,
).scan({
cwd: fullProjectDir,
absolute: true,
})) {
@@ -98,12 +100,12 @@ export namespace Storage {
await Bun.write(dest, JSON.stringify(message))
log.info(`migrating parts for message ${message.id}`)
for await (const partFile of new Bun.Glob(`storage/session/part/${session.id}/${message.id}/*.json`).scan(
{
for await (const partFile of new Bun.Glob(
`storage/session/part/${session.id}/${message.id}/*.json`,
).scan({
cwd: fullProjectDir,
absolute: true,
},
)) {
})) {
const dest = path.join(dir, "part", message.id, path.basename(partFile))
const part = await Bun.file(partFile).json()
log.info("copying", {
@@ -117,6 +119,29 @@ export namespace Storage {
}
}
},
async (dir) => {
for await (const item of new Bun.Glob("session/*/*.json").scan({
cwd: dir,
absolute: true,
})) {
const session = await Bun.file(item).json()
if (!session.projectID) continue
if (!session.summary?.diffs) continue
const { diffs } = session.summary
await Bun.file(path.join(dir, "session_diff", session.id + ".json")).write(
JSON.stringify(diffs),
)
await Bun.file(path.join(dir, "session", session.projectID, session.id + ".json")).write(
JSON.stringify({
...session,
summary: {
additions: diffs.reduce((sum: any, x: any) => sum + x.additions, 0),
deletions: diffs.reduce((sum: any, x: any) => sum + x.deletions, 0),
},
}),
)
}
},
]
const state = lazy(async () => {
@@ -128,9 +153,7 @@ export namespace Storage {
for (let index = migration; index < MIGRATIONS.length; index++) {
log.info("running migration", { index })
const migration = MIGRATIONS[index]
await migration(dir).catch((e) => {
log.error("failed to run migration", { error: e, index })
})
await migration(dir).catch(() => log.error("failed to run migration", { index }))
await Bun.write(path.join(dir, "migration"), (index + 1).toString())
}
return {

View File

@@ -55,6 +55,7 @@ import type {
SessionShareErrors,
SessionDiffData,
SessionDiffResponses,
SessionDiffErrors,
SessionSummarizeData,
SessionSummarizeResponses,
SessionSummarizeErrors,
@@ -475,12 +476,16 @@ class Session extends _HeyApiClient {
}
/**
* Get the diff that resulted from this user message
* Get the diff for this session
*/
public diff<ThrowOnError extends boolean = false>(
options: Options<SessionDiffData, ThrowOnError>,
) {
return (options.client ?? this._client).get<SessionDiffResponses, unknown, ThrowOnError>({
return (options.client ?? this._client).get<
SessionDiffResponses,
SessionDiffErrors,
ThrowOnError
>({
url: "/session/{id}/diff",
...options,
})

View File

@@ -527,7 +527,9 @@ export type Session = {
directory: string
parentID?: string
summary?: {
diffs: Array<FileDiff>
additions: number
deletions: number
diffs?: Array<FileDiff>
}
share?: {
url: string
@@ -1882,6 +1884,9 @@ export type SessionShareResponse = SessionShareResponses[keyof SessionShareRespo
export type SessionDiffData = {
body?: never
path: {
/**
* Session ID
*/
id: string
}
query?: {
@@ -1891,9 +1896,22 @@ export type SessionDiffData = {
url: "/session/{id}/diff"
}
export type SessionDiffErrors = {
/**
* Bad request
*/
400: BadRequestError
/**
* Not found
*/
404: NotFoundError
}
export type SessionDiffError = SessionDiffErrors[keyof SessionDiffErrors]
export type SessionDiffResponses = {
/**
* Successfully retrieved diff
* List of diffs
*/
200: Array<FileDiff>
}