diff --git a/apps/server/src/engine/model.ts b/apps/server/src/engine/model.ts index d7c43330a..5451bed23 100644 --- a/apps/server/src/engine/model.ts +++ b/apps/server/src/engine/model.ts @@ -106,17 +106,7 @@ export async function executeModelTask( "read_workspace", "Read the authorized workspace sources", z.object({ section: z.enum(["mail", "calendar", "files", "all"]) }), - async ({ section }) => { - const w = await service.workspace.snapshot(owner); - return { - mail: section === "mail" || section === "all" ? w.mail : undefined, - events: section === "calendar" || section === "all" ? w.events : undefined, - files: - section === "files" || section === "all" - ? w.files.map(({ url, ...file }) => file) - : undefined, - }; - }, + async ({ section }) => service.workspace.sectionSnapshot(owner, section), ), tool( "read_mail_thread", diff --git a/apps/server/src/workspace.ts b/apps/server/src/workspace.ts index 6d5aafd22..560912052 100644 --- a/apps/server/src/workspace.ts +++ b/apps/server/src/workspace.ts @@ -258,27 +258,42 @@ export class WorkspaceService { await this.db.put(owner, "settings", { id: "google", enabled: true }); await this.db.put(owner, "settings", { id: "seeded", value: true }); } - async snapshot(owner: string, query?: string): Promise { - let mail: Mail[], events: CalendarEvent[]; + private async sourceSections( + owner: string, + wantMail: boolean, + wantEvents: boolean, + query?: string, + ): Promise<{ mail: Mail[]; events: CalendarEvent[]; connected: boolean }> { + let mail: Mail[] = []; + let events: CalendarEvent[] = []; const connected = await this.connected(owner); if (this.config.mode === "live" && connected) { const connection = await this.connection(owner); if (!connection) throw new AppError("Google is disconnected", 409); const google = this.google(owner, connection.id); - [mail, events] = await Promise.all([google.listMail(query), google.listEvents()]); - mail = await this.cacheMail(owner, mail, connection.id); + const [mailResult, eventResult] = await Promise.all([ + wantMail ? google.listMail(query) : Promise.resolve([] as Mail[]), + wantEvents ? google.listEvents() : Promise.resolve([] as CalendarEvent[]), + ]); + mail = mailResult; + events = eventResult; + if (wantMail) mail = await this.cacheMail(owner, mail, connection.id); for (const event of events) await this.db.put(owner, "events", event); } else if (this.config.mode === "sample" && connected) { - mail = await this.db.list(owner, "mail"); - events = await this.db.list(owner, "events"); - if (query) - mail = mail.filter((m) => - `${m.sender} ${m.subject} ${m.body}`.toLowerCase().includes(query.toLowerCase()), - ); - } else { - mail = []; - events = []; + if (wantMail) { + mail = await this.db.list(owner, "mail"); + if (query) + mail = mail.filter((m) => + `${m.sender} ${m.subject} ${m.body}`.toLowerCase().includes(query.toLowerCase()), + ); + } + if (wantEvents) events = await this.db.list(owner, "events"); } + return { mail, events, connected }; + } + + async snapshot(owner: string, query?: string): Promise { + const { mail, events, connected } = await this.sourceSections(owner, true, true, query); const tokens = this.config.mode === "live" ? await this.googleAuth.tokens(owner) : null; return { mode: this.config.mode, @@ -327,6 +342,27 @@ export class WorkspaceService { }, }; } + + async sectionSnapshot( + owner: string, + section: "mail" | "calendar" | "files" | "all", + ): Promise<{ mail?: Mail[]; events?: CalendarEvent[]; files?: Omit[] }> { + const wantMail = section === "mail" || section === "all"; + const wantEvents = section === "calendar" || section === "all"; + const wantFiles = section === "files" || section === "all"; + const { mail, events } = + wantMail || wantEvents + ? await this.sourceSections(owner, wantMail, wantEvents) + : { mail: [], events: [] }; + const files = wantFiles + ? (await this.files.list(owner)).map(({ url: _url, ...file }) => file) + : undefined; + return { + mail: wantMail ? mail.sort((a, b) => b.date.localeCompare(a.date)) : undefined, + events: wantEvents ? events.sort((a, b) => a.start.localeCompare(b.start)) : undefined, + files, + }; + } async prepare(owner: string, input: ProposalInput, connectionId?: string) { if (input.kind === "email.send") { for (const id of input.data.attachmentIds) await this.files.get(owner, id); diff --git a/tests/persistence.test.ts b/tests/persistence.test.ts index e371655cb..88f10d5b2 100644 --- a/tests/persistence.test.ts +++ b/tests/persistence.test.ts @@ -31,3 +31,54 @@ test("idle Postgres client errors are logged instead of crashing the process", a await pool.end(); } }); +test("workspace section reads match the selected snapshot sections", async () => { + const { WorkspaceService } = await import("../apps/server/src/workspace.ts"); + const { Files } = await import("../apps/server/src/files.ts"); + const store = await createStore(); + try { + const config = { mode: "sample" } as unknown as import("../apps/server/src/config.ts").Config; + const workspace = new WorkspaceService( + store, + config, + new Files(store, config, { sign: () => "sig" } as never), + {} as never, + ); + await store.put("owner", "settings", { id: "google", enabled: true }); + await store.put("owner", "mail", { + id: "m1", + sender: "a@example.com", + subject: "Subject", + body: "Body", + date: "2026-01-02T00:00:00Z", + label: "INBOX", + threadId: "t1", + attachments: [], + }); + await store.put("owner", "events", { + id: "e1", + title: "Event", + start: "2026-01-02T00:00:00Z", + end: "2026-01-02T01:00:00Z", + }); + const full = await workspace.snapshot("owner"); + + const mail = await workspace.sectionSnapshot("owner", "mail"); + assert.deepEqual(mail.mail, full.mail); + assert.equal(mail.events, undefined); + assert.equal(mail.files, undefined); + + const calendar = await workspace.sectionSnapshot("owner", "calendar"); + assert.deepEqual(calendar.events, full.events); + assert.equal(calendar.mail, undefined); + + const files = await workspace.sectionSnapshot("owner", "files"); + assert.deepEqual( + files.files, + full.files.map(({ url: _url, ...file }) => file), + ); + assert.equal(files.mail, undefined); + assert.equal(files.events, undefined); + } finally { + await store.close(); + } +});