diff --git a/src/lib/mcp/__tests__/record-candidates-budget.test.ts b/src/lib/mcp/__tests__/record-candidates-budget.test.ts index 9f8332f1b..26e74fbdf 100644 --- a/src/lib/mcp/__tests__/record-candidates-budget.test.ts +++ b/src/lib/mcp/__tests__/record-candidates-budget.test.ts @@ -80,4 +80,32 @@ describe("loadRecordCandidates under a budget", () => { expect(out).toHaveLength(10); expect(truncated).toEqual([]); }); + + it("does not call a kind holding exactly its cap truncated", async () => { + prismaMock.encounter.findMany.mockResolvedValue(encounters(5)); + const truncated: string[] = []; + const out = await loadRecordCandidates("u1", "visit", "", { + kinds: new Set(["visit"]), + rowsPerKind: 5, + onTruncated: (kind) => truncated.push(kind), + }); + expect(out).toHaveLength(5); + expect(truncated).toEqual([]); + // One row past the cap is asked for, as the probe. + expect(prismaMock.encounter.findMany).toHaveBeenCalledWith( + expect.objectContaining({ take: 6 }), + ); + }); + + it("calls a kind with a row past its cap truncated, and keeps the cap", async () => { + prismaMock.encounter.findMany.mockResolvedValue(encounters(6)); + const truncated: string[] = []; + const out = await loadRecordCandidates("u1", "visit", "", { + kinds: new Set(["visit"]), + rowsPerKind: 5, + onTruncated: (kind) => truncated.push(kind), + }); + expect(out).toHaveLength(5); + expect(truncated).toEqual(["visit"]); + }); }); diff --git a/src/lib/mcp/record-search.ts b/src/lib/mcp/record-search.ts index 75f146bc9..79cca5ebd 100644 --- a/src/lib/mcp/record-search.ts +++ b/src/lib/mcp/record-search.ts @@ -361,66 +361,70 @@ async function visitCandidates( ...(empty ? { occurredAt: { lte: new Date() } } : {}), }, orderBy: [{ occurredAt: "desc" }, { id: "asc" }], - take: read.take, + take: probeTake(read), include: { practitioner: PRACTITIONER_SELECT }, }); - return mapUntilStopped(rows, read, (row): SearchCandidate => { - const practitioner = livePractitioner(row.practitioner); - const reason = decryptText(row.reasonEncrypted); - const outcome = decryptText(row.outcomeEncrypted); - const bodySite = decryptText(row.bodySiteEncrypted); - const day = userDayKey(row.occurredAt, owner.timezone); - const kindLabel = encounterKindLabel(row.kind, owner.display); - const title = reason?.replace(/\s+/g, " ").trim() || null; - return { - id: `visit:${row.id}`, - title: visitTitle({ - kind: row.kind, - day, - bodySite: bodySite?.replace(/\s+/g, " ").trim() ?? null, - laterality: row.laterality, - practitioner: practitioner?.name ?? null, - status: row.status, - }), - url: recordUrl(origin, "visit", row.id), - kindOrder: KIND_ORDER.visit, - date: row.occurredAt.getTime(), - record: true, - hit: { - kind: "visit", - recordId: row.id, - title: title ?? kindLabel, - context: compact([ - title ? kindLabel : null, - practitioner?.name, - sitePhrase(bodySite, row.laterality, owner.display), - ]), - day, - path: recordPath("visit", row.id), - }, - fields: [ - { text: bodySite ?? "", weight: WEIGHT.primary }, - { text: reason ?? "", weight: WEIGHT.primary, body: true }, - { text: practitioner?.name ?? "", weight: WEIGHT.primary }, - { text: practitioner?.specialty ?? "", weight: WEIGHT.primary }, - { - text: [ - "visit", - enumWords(row.kind), - encounterKindLabel(row.kind, "en"), - encounterKindLabel(row.kind, owner.locale), - ENCOUNTER_KIND_ALIASES[row.kind] ?? "", - ].join(" "), - weight: WEIGHT.secondary, + return mapUntilStopped( + withinCap(rows, read), + read, + (row): SearchCandidate => { + const practitioner = livePractitioner(row.practitioner); + const reason = decryptText(row.reasonEncrypted); + const outcome = decryptText(row.outcomeEncrypted); + const bodySite = decryptText(row.bodySiteEncrypted); + const day = userDayKey(row.occurredAt, owner.timezone); + const kindLabel = encounterKindLabel(row.kind, owner.display); + const title = reason?.replace(/\s+/g, " ").trim() || null; + return { + id: `visit:${row.id}`, + title: visitTitle({ + kind: row.kind, + day, + bodySite: bodySite?.replace(/\s+/g, " ").trim() ?? null, + laterality: row.laterality, + practitioner: practitioner?.name ?? null, + status: row.status, + }), + url: recordUrl(origin, "visit", row.id), + kindOrder: KIND_ORDER.visit, + date: row.occurredAt.getTime(), + record: true, + hit: { + kind: "visit", + recordId: row.id, + title: title ?? kindLabel, + context: compact([ + title ? kindLabel : null, + practitioner?.name, + sitePhrase(bodySite, row.laterality, owner.display), + ]), + day, + path: recordPath("visit", row.id), }, - { - text: sideTerms(row.laterality, owner.locale) ?? "", - weight: WEIGHT.secondary, - }, - { text: outcome ?? "", weight: WEIGHT.secondary, body: true }, - ], - }; - }); + fields: [ + { text: bodySite ?? "", weight: WEIGHT.primary }, + { text: reason ?? "", weight: WEIGHT.primary, body: true }, + { text: practitioner?.name ?? "", weight: WEIGHT.primary }, + { text: practitioner?.specialty ?? "", weight: WEIGHT.primary }, + { + text: [ + "visit", + enumWords(row.kind), + encounterKindLabel(row.kind, "en"), + encounterKindLabel(row.kind, owner.locale), + ENCOUNTER_KIND_ALIASES[row.kind] ?? "", + ].join(" "), + weight: WEIGHT.secondary, + }, + { + text: sideTerms(row.laterality, owner.locale) ?? "", + weight: WEIGHT.secondary, + }, + { text: outcome ?? "", weight: WEIGHT.secondary, body: true }, + ], + }; + }, + ); } /** @@ -481,53 +485,57 @@ async function conditionCandidates( const rows = await prisma.illnessEpisode.findMany({ where: { userId, deletedAt: null }, orderBy: [{ onsetAt: "desc" }, { id: "asc" }], - take: read.take, + take: probeTake(read), select: CONDITION_SELECT, }); - return mapUntilStopped(rows, read, (row): SearchCandidate => { - const bodySite = decryptText(row.bodySiteEncrypted); - const day = userDayKey(row.onsetAt, owner.timezone); - const site = bodySite?.replace(/\s+/g, " ").trim(); - return { - id: `condition:${row.id}`, - title: scrubFenceMarkers( - `${row.label}${site ? ` · ${site}` : ""} (${illnessTypeLabel(row.type, "en")}, since ${day})`, - ), - url: recordUrl(origin, "condition", row.id), - kindOrder: KIND_ORDER.condition, - date: row.onsetAt.getTime(), - record: true, - hit: { - kind: "condition", - recordId: row.id, - title: row.label, - context: compact([ - illnessTypeLabel(row.type, owner.display), - sitePhrase(bodySite, row.laterality, owner.display), - ]), - day, - path: recordPath("condition", row.id), - }, - fields: [ - { text: row.label, weight: WEIGHT.primary }, - { text: bodySite ?? "", weight: WEIGHT.primary }, - { - text: [ - "condition illness", - enumWords(row.type), - illnessTypeLabel(row.type, "en"), - illnessTypeLabel(row.type, owner.locale), - enumWords(row.lifecycle), - ].join(" "), - weight: WEIGHT.secondary, + return mapUntilStopped( + withinCap(rows, read), + read, + (row): SearchCandidate => { + const bodySite = decryptText(row.bodySiteEncrypted); + const day = userDayKey(row.onsetAt, owner.timezone); + const site = bodySite?.replace(/\s+/g, " ").trim(); + return { + id: `condition:${row.id}`, + title: scrubFenceMarkers( + `${row.label}${site ? ` · ${site}` : ""} (${illnessTypeLabel(row.type, "en")}, since ${day})`, + ), + url: recordUrl(origin, "condition", row.id), + kindOrder: KIND_ORDER.condition, + date: row.onsetAt.getTime(), + record: true, + hit: { + kind: "condition", + recordId: row.id, + title: row.label, + context: compact([ + illnessTypeLabel(row.type, owner.display), + sitePhrase(bodySite, row.laterality, owner.display), + ]), + day, + path: recordPath("condition", row.id), }, - { - text: sideTerms(row.laterality, owner.locale) ?? "", - weight: WEIGHT.secondary, - }, - ], - }; - }); + fields: [ + { text: row.label, weight: WEIGHT.primary }, + { text: bodySite ?? "", weight: WEIGHT.primary }, + { + text: [ + "condition illness", + enumWords(row.type), + illnessTypeLabel(row.type, "en"), + illnessTypeLabel(row.type, owner.locale), + enumWords(row.lifecycle), + ].join(" "), + weight: WEIGHT.secondary, + }, + { + text: sideTerms(row.laterality, owner.locale) ?? "", + weight: WEIGHT.secondary, + }, + ], + }; + }, + ); } const DOCUMENT_SELECT = { @@ -610,12 +618,12 @@ async function documentCandidates( includeAiDeferred: boolean, ): Promise { const nameTerms = documentNameTerms(query, words); - const [rows, content, named] = await Promise.all([ + const [probed, content, named] = await Promise.all([ prisma.inboundDocument.findMany({ where: { userId, deletedAt: null }, select: DOCUMENT_SELECT, orderBy: [{ createdAt: "desc" }, { id: "asc" }], - take: read.take, + take: probeTake(read), }), words.length > 0 ? contentMatches(userId, words, includeAiDeferred) @@ -638,6 +646,9 @@ async function documentCandidates( }) : Promise.resolve([] as Array<{ id: string }>), ]); + // The cap applies to the newest read; the extra rows below are content and + // name hits beyond it, already bounded by their own read. + const rows = withinCap(probed, read); // A content or name hit on a document older than the newest `take` still // counts; the extra read is bounded like the others. const seen = new Set(rows.map((r) => r.id)); @@ -654,58 +665,51 @@ async function documentCandidates( }) : []; - // The cap applies to the newest read; the extra rows are content and name - // hits beyond it, already bounded by their own read. - if (rows.length >= read.take) read.onTruncated(); - return mapUntilStopped( - [...rows, ...extra], - { ...read, take: Infinity }, - (row): SearchCandidate => { - const dated = row.documentDate ?? row.reportDate; - const kindLabel = documentKindLabel(row.kind, "en"); - const name = row.title ?? row.filename ?? kindLabel; - return { - id: `document:${row.id}`, - title: scrubFenceMarkers( - `${name} (${kindLabel}${dated ? `, ${dateOnly(dated)}` : ""})`, - ), - url: recordUrl(origin, "document", row.id), - kindOrder: KIND_ORDER.document, - date: (dated ?? row.createdAt).getTime(), - record: true, - hit: { - kind: "document", - recordId: row.id, - title: - row.title ?? - row.filename ?? - documentKindLabel(row.kind, owner.display), - context: compact([ - documentKindLabel(row.kind, owner.display), - row.title ? row.filename : null, - ]), - day: dated - ? dateOnly(dated) - : userDayKey(row.createdAt, owner.timezone), - path: recordPath("document", row.id), + return mapUntilStopped([...rows, ...extra], read, (row): SearchCandidate => { + const dated = row.documentDate ?? row.reportDate; + const kindLabel = documentKindLabel(row.kind, "en"); + const name = row.title ?? row.filename ?? kindLabel; + return { + id: `document:${row.id}`, + title: scrubFenceMarkers( + `${name} (${kindLabel}${dated ? `, ${dateOnly(dated)}` : ""})`, + ), + url: recordUrl(origin, "document", row.id), + kindOrder: KIND_ORDER.document, + date: (dated ?? row.createdAt).getTime(), + record: true, + hit: { + kind: "document", + recordId: row.id, + title: + row.title ?? + row.filename ?? + documentKindLabel(row.kind, owner.display), + context: compact([ + documentKindLabel(row.kind, owner.display), + row.title ? row.filename : null, + ]), + day: dated + ? dateOnly(dated) + : userDayKey(row.createdAt, owner.timezone), + path: recordPath("document", row.id), + }, + contentWords: content.get(row.id), + fields: [ + { text: row.title ?? "", weight: WEIGHT.primary }, + { text: row.filename ?? "", weight: WEIGHT.primary }, + { + text: [ + "document", + enumWords(row.kind), + kindLabel, + documentKindLabel(row.kind, owner.locale), + ].join(" "), + weight: WEIGHT.secondary, }, - contentWords: content.get(row.id), - fields: [ - { text: row.title ?? "", weight: WEIGHT.primary }, - { text: row.filename ?? "", weight: WEIGHT.primary }, - { - text: [ - "document", - enumWords(row.kind), - kindLabel, - documentKindLabel(row.kind, owner.locale), - ].join(" "), - weight: WEIGHT.secondary, - }, - ], - }; - }, - ); + ], + }; + }); } async function vaccinationCandidates( @@ -717,68 +721,72 @@ async function vaccinationCandidates( const rows = await prisma.vaccinationRecord.findMany({ where: { userId, deletedAt: null }, orderBy: [{ occurredAt: "desc" }, { id: "asc" }], - take: read.take, + take: probeTake(read), select: VACCINATION_SELECT, }); - return mapUntilStopped(rows, read, (row): SearchCandidate => { - const practitioner = livePractitioner(row.practitioner); - const entry = resolveVaccineEntry({ antigenSlug: row.antigenSlug }); - const own = ownVaccine(row); - const catalogEn = vaccineCatalogName(row.antigenSlug, "en"); - const name = - row.vaccineName ?? - catalogEn ?? - own?.name ?? - row.antigenSlug ?? - "Vaccination"; - const dose = row.doseNumber ? `dose ${row.doseNumber}, ` : ""; - const day = dateOnly(row.occurredAt); - return { - id: `vaccination:${row.id}`, - title: scrubFenceMarkers(`${name} (${dose}${day})`), - url: recordUrl(origin, "vaccination", row.id), - kindOrder: KIND_ORDER.vaccination, - date: row.occurredAt.getTime(), - record: true, - hit: { - kind: "vaccination", - recordId: row.id, - title: - row.vaccineName ?? - vaccineCatalogName(row.antigenSlug, owner.display) ?? - own?.name ?? - row.antigenSlug ?? - null, - context: compact([practitioner?.name]), - day, - path: recordPath("vaccination", row.id), - }, - fields: [ - { text: row.vaccineName ?? "", weight: WEIGHT.primary }, - { - text: [ - catalogEn ?? "", - vaccineCatalogName(row.antigenSlug, owner.locale) ?? "", - row.antigenSlug - ? enumWords(row.antigenSlug.replace(/-/g, "_")) - : "", - ...(entry?.synonyms ?? []), - own?.name ?? "", - ...(own?.components ?? []).map((antigen) => - enumWords(antigen.replace(/-/g, "_")), - ), - ].join(" "), - weight: WEIGHT.primary, + return mapUntilStopped( + withinCap(rows, read), + read, + (row): SearchCandidate => { + const practitioner = livePractitioner(row.practitioner); + const entry = resolveVaccineEntry({ antigenSlug: row.antigenSlug }); + const own = ownVaccine(row); + const catalogEn = vaccineCatalogName(row.antigenSlug, "en"); + const name = + row.vaccineName ?? + catalogEn ?? + own?.name ?? + row.antigenSlug ?? + "Vaccination"; + const dose = row.doseNumber ? `dose ${row.doseNumber}, ` : ""; + const day = dateOnly(row.occurredAt); + return { + id: `vaccination:${row.id}`, + title: scrubFenceMarkers(`${name} (${dose}${day})`), + url: recordUrl(origin, "vaccination", row.id), + kindOrder: KIND_ORDER.vaccination, + date: row.occurredAt.getTime(), + record: true, + hit: { + kind: "vaccination", + recordId: row.id, + title: + row.vaccineName ?? + vaccineCatalogName(row.antigenSlug, owner.display) ?? + own?.name ?? + row.antigenSlug ?? + null, + context: compact([practitioner?.name]), + day, + path: recordPath("vaccination", row.id), }, - { - text: "vaccination vaccine immunization impfung", - weight: WEIGHT.secondary, - }, - { text: practitioner?.name ?? "", weight: WEIGHT.secondary }, - { text: practitioner?.specialty ?? "", weight: WEIGHT.secondary }, - ], - }; - }); + fields: [ + { text: row.vaccineName ?? "", weight: WEIGHT.primary }, + { + text: [ + catalogEn ?? "", + vaccineCatalogName(row.antigenSlug, owner.locale) ?? "", + row.antigenSlug + ? enumWords(row.antigenSlug.replace(/-/g, "_")) + : "", + ...(entry?.synonyms ?? []), + own?.name ?? "", + ...(own?.components ?? []).map((antigen) => + enumWords(antigen.replace(/-/g, "_")), + ), + ].join(" "), + weight: WEIGHT.primary, + }, + { + text: "vaccination vaccine immunization impfung", + weight: WEIGHT.secondary, + }, + { text: practitioner?.name ?? "", weight: WEIGHT.secondary }, + { text: practitioner?.specialty ?? "", weight: WEIGHT.secondary }, + ], + }; + }, + ); } export interface RecordSearchOptions { @@ -818,6 +826,22 @@ interface KindRead { onTruncated: () => void; } +/** + * A kind's newest read asks for one row past its cap. The extra row only + * answers whether anything older is left: a kind holding exactly `take` rows + * was read whole and is not reported cut short. + */ +function probeTake(read: KindRead): number { + return read.take + 1; +} + +/** The rows within the cap, telling `onTruncated` when the probe row came. */ +function withinCap(rows: Row[], read: KindRead): Row[] { + if (rows.length <= read.take) return rows; + read.onTruncated(); + return rows.slice(0, read.take); +} + /** * Map a kind's rows newest first until `shouldStop` says stop. The rows are * read already; what this bounds is the decrypting and the building, which is @@ -828,7 +852,6 @@ function mapUntilStopped( read: KindRead, build: (row: Row) => Out, ): Out[] { - if (rows.length >= read.take) read.onTruncated(); const out: Out[] = []; for (const row of rows) { if (read.shouldStop()) { diff --git a/src/lib/search/__tests__/search-records.test.ts b/src/lib/search/__tests__/search-records.test.ts index 0b9574224..deacfd9a2 100644 --- a/src/lib/search/__tests__/search-records.test.ts +++ b/src/lib/search/__tests__/search-records.test.ts @@ -79,7 +79,24 @@ describe("scanNewest", () => { const s = source(10_000, 1_000_000); await scanNewest({ budget: b, wanted: 3, batch: 100, maxRows: 450, ...s }); expect(b.rows).toBe(450); - expect(s.reads).toEqual([100, 100, 100, 100, 50]); + // The last batch probes one row past the cap and does not keep it. + expect(s.reads).toEqual([100, 100, 100, 100, 51]); + expect(b.truncated).toBe(true); + }); + + it("does not call a source holding exactly the cap truncated", async () => { + const b = budget(); + const s = source(450, 1_000_000); + await scanNewest({ budget: b, wanted: 3, batch: 100, maxRows: 450, ...s }); + expect(b.rows).toBe(450); + expect(b.truncated).toBe(false); + }); + + it("calls a source one row past the cap truncated", async () => { + const b = budget(); + const s = source(451, 1_000_000); + await scanNewest({ budget: b, wanted: 3, batch: 100, maxRows: 450, ...s }); + expect(b.rows).toBe(450); expect(b.truncated).toBe(true); }); diff --git a/src/lib/search/search-records.ts b/src/lib/search/search-records.ts index 189d665c1..152ac3f8d 100644 --- a/src/lib/search/search-records.ts +++ b/src/lib/search/search-records.ts @@ -143,22 +143,30 @@ export async function scanNewest(args: { budget.truncated = true; break; } - const take = Math.min(batch, maxRows - scanned); - const rows = await read(cursor, take); + const remaining = maxRows - scanned; + // The batch that reaches the cap asks for one row more than it keeps: + // that row only answers whether anything is left, so a source holding + // exactly `maxRows` rows is read whole and not reported cut short. + const last = remaining <= batch; + const take = last ? remaining : batch; + const fetched = await read(cursor, last ? take + 1 : take); + const more = last && fetched.length > take; + const rows = more ? fetched.slice(0, take) : fetched; scanned += rows.length; budget.rows += rows.length; for (const row of rows) { const candidate = match(row); if (candidate) out.push(candidate); } - // A short batch is the end of the source. - if (rows.length < take) break; - // One more than the limit: the caller learns there is more. - if (out.length > wanted) break; - if (scanned >= maxRows) { + if (more) { budget.truncated = true; break; } + // A short batch, or the cap reached with nothing beyond it, is the end + // of the source. + if (rows.length < take || last) break; + // One more than the limit: the caller learns there is more. + if (out.length > wanted) break; cursor = rows[rows.length - 1].id; } return out;