603 lines · 19.3 KB · TypeScript
RawDownload
import { describe, expect, it } from 'vitest' import {  ANNOTATION_LIMITS,  annotationPayloadText,  annotationRequiresImageInput,  imageAnnotationSourceIsFixed,  pdfAnnotationSourceIsFixed,  parseSideChatAnnotationText,  prepareAnnotationsForAgent,  sanitizeAnnotations,  sideChatAnnotationText,  validateAnnotations,  type Annotation,  type ImagePointAnnotation,  type PdfAnnotation,  type TextAnnotation} from './annotations' const textAnnotation = (overrides: Partial<TextAnnotation> = {}): TextAnnotation => ({  id: 'annotation-1',  kind: 'text',  target: 'agent',  quote: 'The confidence intervals overlap.',  source: {    kind: 'agent-message',    sessionId: 'session-1',    messageId: 'message-1'  },  ...overrides}) const imageAnnotation = (): ImagePointAnnotation => ({  id: 'image-1',  kind: 'image-point',  target: 'agent',  note: 'Inspect the peak.',  source: {    kind: 'artifact-version',    projectId: 'project-1',    sessionId: 'session-1',    versionId: 'version-1',    name: 'figure.png',    path: 'artifact-version:project-1/session-1/artifact-1/version-1',    mimeType: 'image/png'  },  point: { x: 0.5, y: 0.25 },  naturalSize: { width: 1000, height: 800 }}) const pdfAnnotation = (overrides: Partial<PdfAnnotation> = {}): PdfAnnotation => ({  id: 'pdf-1',  kind: 'pdf',  target: 'agent',  source: {    kind: 'upload-version',    projectId: 'project-1',    sessionId: 'session-1',    versionId: 'version-7',    name: 'paper.pdf',    path: 'upload-version:project-1/session-1/version-7',    checksum: 'a'.repeat(64)  },  selector: {    kind: 'text',    pageNumber: 4,    exact: 'The confidence intervals overlap.',    prefix: 'Result: ',    suffix: ' This matters.',    position: { start: 8, end: 41 },    quads: [{ x: 0.1, y: 0.2, width: 0.4, height: 0.03 }],    extractorVersion: 'pdfjs-5.4.624'  },  ...overrides}) const pdfRegionAnnotation = (): PdfAnnotation => ({  ...pdfAnnotation(),  id: 'pdf-region-1',  selector: {    kind: 'region',    pageNumber: 6,    rect: { x: 0.2, y: 0.3, width: 0.4, height: 0.25 },    pageRotation: 0,    text: 'Figure 2. Retrieval evaluator.',    image: { mimeType: 'image/png', data: 'AQID', byteLength: 3 }  }}) describe('ordinary text source identity in Agent input', () => {  it.each(['file', 'message'] as const)(    'distinguishes identical quotes from different %s sources',    (kind) => {      const from = (label: string): TextAnnotation =>        textAnnotation({          quote: 'The value is 42.',          note: 'Correct the referenced source.',          source:            kind === 'file'              ? {                  kind: 'project-file',                  projectId: 'project-1',                  sessionId: 'session-1',                  name: `file-${label}.md`,                  versionId: `version-${label}`,                  path: `artifact-version:project-1/session-1/file-${label}/version-${label}`                }              : { kind: 'agent-message', sessionId: 'session-1', messageId: `message-${label}` }        })      const a = from('A')      const b = from('B')      expect(validateAnnotations([a])).toBeUndefined()      expect(validateAnnotations([b])).toBeUndefined()      const promptA = prepareAnnotationsForAgent('', [a])      const promptB = prepareAnnotationsForAgent('', [b])      expect(promptA.promptText).toContain(a.quote)      expect(promptB.promptText).toContain(b.quote)      expect(promptB).not.toEqual(promptA)    }  )}) describe('annotations', () => {  it('restores bounded historical annotations even when new source metadata exceeds the send budget', () => {    const historical = Array.from({ length: 10 }, (_, index) =>      textAnnotation({        id: `history-${index}`,        quote: 'x'.repeat(940)      })    )    expect(validateAnnotations(historical)).toBe('payload-too-large')    expect(sanitizeAnnotations(historical)).toEqual(historical)  })   it('round-trips text anchors and rejects malformed anchor positions', () => {    const selected = textAnnotation({      quote: 'repeat',      anchor: {        position: { start: 12, end: 18 },        prefix: 'repeat then '      }    })    expect(sanitizeAnnotations(JSON.parse(JSON.stringify([selected])))).toEqual([selected])    expect(      sanitizeAnnotations([{ ...selected, anchor: { position: { start: 12, end: 17 } } }])    ).toEqual([])    expect(      sanitizeAnnotations([{ ...selected, anchor: { position: { start: -1, end: 5 } } }])    ).toEqual([])  })   it('prepares omitted PDF bitmaps explicitly and numbers only actual image attachments', () => {    const region = pdfRegionAnnotation()    const omitted: PdfAnnotation = {      ...region,      id: 'omitted',      selector: {        kind: 'region',        pageNumber: 6,        rect: { x: 0.2, y: 0.3, width: 0.4, height: 0.25 },        pageRotation: 0,        text: 'Figure 2.',        imageOmissionReason: 'session-budget'      }    }    expect(sanitizeAnnotations(JSON.parse(JSON.stringify([omitted])))).toEqual([omitted])    expect(annotationRequiresImageInput(omitted)).toBe(false)    const prepared = prepareAnnotationsForAgent('', [omitted, region])    expect(prepared.images).toEqual([{ mimeType: 'image/png', data: 'AQID', byteLength: 3 }])    const payload = JSON.parse(prepared.promptText.slice('[Annotations]\n'.length))    expect(payload.items[0]).toMatchObject({      type: 'pdf-region',      imageOmissionReason: 'session-budget',      source: { versionId: region.source.versionId, page: 6 }    })    expect(payload.items[0]).not.toHaveProperty('image')    expect(payload.items[1]).toMatchObject({ image: 1 })    expect(payload.items[1]).not.toHaveProperty('imageOmissionReason')    expect(      sanitizeAnnotations([        { ...omitted, selector: { ...omitted.selector, imageOmissionReason: 'unknown' } }      ])    ).toEqual([])    expect(      sanitizeAnnotations([        { ...omitted, selector: { ...omitted.selector, imageOmissionReason: undefined } }      ])    ).toEqual([])    expect(      sanitizeAnnotations([        { ...region, selector: { ...region.selector, imageOmissionReason: 'session-budget' } }      ])    ).toEqual([])  })   it('includes deduplicated immutable file references for ordinary text quotes', () => {    const annotation = textAnnotation({      source: {        kind: 'project-file',        projectId: 'project-1',        sessionId: 'session-1',        name: 'results.md',        versionId: 'version-1',        path: 'artifact-version:project-1/session-1/artifact-1/version-1'      }    })    const prepared = prepareAnnotationsForAgent('', [annotation, { ...annotation, id: 'another' }])    expect(prepared.referencedArtifacts).toEqual([      {        id: 'artifact-1',        sourceFileId: 'artifact-1',        name: 'results.md',        source: 'artifact',        versionId: 'version-1',        path: 'artifact-version:project-1/session-1/artifact-1/version-1'      }    ])    expect(      prepareAnnotationsForAgent('', [annotation], prepared.referencedArtifacts).referencedArtifacts    ).toEqual(prepared.referencedArtifacts)  })   it('still parses historical Side chat quotes without source metadata', () => {    expect(      parseSideChatAnnotationText(        '[Annotations]\n{"items":[{"type":"quote","content":"Evidence"}]}'      )    ).toEqual({      text: '',      items: [{ type: 'quote', content: 'Evidence' }]    })  })   it('projects mixed annotations into canonical Side chat text without an image attachment', () => {    const annotations: Annotation[] = [      textAnnotation({ note: 'Explain this caveat.' }),      imageAnnotation()    ]    const expected =      'Compare these observations.\n\n[Annotations]\n' +      '{"items":[{"type":"quote","content":"The confidence intervals overlap.","source":{"kind":"agent-message","sessionId":"session-1","messageId":"message-1"},"instruction":"Explain this caveat."},{"type":"image-point","source":{"kind":"artifact-version","artifactId":"artifact-1","versionId":"version-1","name":"figure.png"},"x":500,"y":200,"instruction":"Inspect the peak."}]}'     expect(sideChatAnnotationText('  Compare these observations.  ', annotations)).toBe(expected)    expect(expected).not.toContain('imageAttachment')    expect(parseSideChatAnnotationText(expected)).toEqual({      text: 'Compare these observations.',      items: [        {          type: 'quote',          content: 'The confidence intervals overlap.',          source: { kind: 'agent-message', sessionId: 'session-1', messageId: 'message-1' },          instruction: 'Explain this caveat.'        },        {          type: 'image-point',          source: {            kind: 'artifact-version',            artifactId: 'artifact-1',            versionId: 'version-1',            name: 'figure.png'          },          x: 500,          y: 200,          instruction: 'Inspect the peak.'        }      ]    })  })   it('projects and parses an annotation-only Side chat message', () => {    const projected = sideChatAnnotationText('', [textAnnotation()])     expect(projected).toBe(      '[Annotations]\n{"items":[{"type":"quote","content":"The confidence intervals overlap.","source":{"kind":"agent-message","sessionId":"session-1","messageId":"message-1"}}]}'    )    expect(parseSideChatAnnotationText(projected)).toEqual({      text: '',      items: [        {          type: 'quote',          content: 'The confidence intervals overlap.',          source: { kind: 'agent-message', sessionId: 'session-1', messageId: 'message-1' }        }      ]    })  })   it.each([    ['malformed JSON', '[Annotations]\n{"items":'],    [      'unknown payload field',      '[Annotations]\n{"items":[{"type":"quote","content":"Evidence","future":true}]}'    ],    ['unknown item type', '[Annotations]\n{"items":[{"type":"future","content":"Evidence"}]}'],    ['trailing text', '[Annotations]\n{"items":[{"type":"quote","content":"Evidence"}]}\ntrailing'],    [      'noncanonical JSON whitespace',      '[Annotations]\n{ "items": [{"type":"quote","content":"Evidence"}]}'    ],    ['empty items', '[Annotations]\n{"items":[]}'],    [      'excessive quote',      `[Annotations]\n${JSON.stringify({        items: [{ type: 'quote', content: 'x'.repeat(ANNOTATION_LIMITS.quote + 1) }]      })}`    ]  ])('treats %s as ordinary Side chat text', (_reason, value) => {    expect(parseSideChatAnnotationText(value)).toBeUndefined()  })   it('uses the complete EOF suffix when the user text contains the annotation marker', () => {    const projected = sideChatAnnotationText('Explain the literal marker [Annotations]\n here.', [      textAnnotation()    ])     expect(parseSideChatAnnotationText(projected)).toEqual({      text: 'Explain the literal marker [Annotations]\n here.',      items: [        {          type: 'quote',          content: 'The confidence intervals overlap.',          source: { kind: 'agent-message', sessionId: 'session-1', messageId: 'message-1' }        }      ]    })  })   it('serializes text annotations as bounded structured Agent context', () => {    expect(annotationPayloadText([textAnnotation({ note: 'Explain this caveat.' })])).toBe(      '[Annotations]\n' +        JSON.stringify({          items: [            {              type: 'quote',              content: 'The confidence intervals overlap.',              source: { kind: 'agent-message', sessionId: 'session-1', messageId: 'message-1' },              instruction: 'Explain this caveat.'            }          ]        })    )  })   it('preserves complete managed project-file identity and rejects partial identity', () => {    const annotation = textAnnotation({      source: {        kind: 'project-file',        projectId: 'project-1',        sessionId: 'session-1',        path: '/stale/report.md',        name: 'report.md',        fileSource: 'artifact',        sourceFileId: 'artifact-1',        versionId: 'version-2'      }    })     expect(sanitizeAnnotations([annotation])).toEqual([annotation])    expect(      sanitizeAnnotations([        {          ...annotation,          source: { ...annotation.source, sourceFileId: undefined }        }      ])    ).toEqual([])  })   it('rejects managed project-file identity without an immutable or consistent Version', () => {    const annotation = textAnnotation({      source: {        kind: 'project-file',        projectId: 'project-1',        sessionId: 'session-1',        path: '/stale/report.md',        name: 'report.md',        fileSource: 'artifact',        sourceFileId: 'artifact-1'      }    })     expect(sanitizeAnnotations([annotation])).toEqual([])    expect(      sanitizeAnnotations([        {          ...annotation,          source: {            ...annotation.source,            path: 'artifact-version:project-1/session-1/artifact-2/version-2',            versionId: 'version-2'          }        }      ])    ).toEqual([])  })   it('keeps immutable PDF text anchors in persisted and Agent context', () => {    const annotation = pdfAnnotation()     const sanitized = sanitizeAnnotations([annotation])     expect(sanitized).toEqual([annotation])    expect(annotationPayloadText(sanitized)).toBe(      '[Annotations]\n' +        JSON.stringify({          items: [            {              type: 'quote',              content: 'The confidence intervals overlap.',              source: {                kind: 'pdf',                versionId: 'version-7',                name: 'paper.pdf',                checksum: 'a'.repeat(64),                page: 4,                selector: {                  prefix: 'Result: ',                  suffix: ' This matters.',                  position: { start: 8, end: 41 },                  quads: [{ x: 0.1, y: 0.2, width: 0.4, height: 0.03 }],                  extractorVersion: 'pdfjs-5.4.624'                }              }            }          ]        })    )    expect(parseSideChatAnnotationText(sideChatAnnotationText('', sanitized))).toEqual({      text: '',      items: [        {          type: 'quote',          content: 'The confidence intervals overlap.',          source: expect.objectContaining({            kind: 'pdf',            versionId: 'version-7',            checksum: 'a'.repeat(64),            page: 4          })        }      ]    })    expect(      sanitizeAnnotations([        {          ...annotation,          selector: { ...annotation.selector, pageNumber: 0 }        }      ])    ).toEqual([])  })   it('keeps a Literature PDF anchor without inventing a source Session', () => {    const annotation = pdfAnnotation({      source: {        kind: 'literature-attachment-version',        projectId: 'project-1',        versionId: 'attachment-version-1',        name: 'library-paper.pdf',        path: 'literature-attachment-version:attachment-version-1',        checksum: 'b'.repeat(64)      }    })     expect(sanitizeAnnotations([annotation])).toEqual([annotation])    expect(pdfAnnotationSourceIsFixed(annotation.source)).toBe(true)  })   it('keeps a bounded PDF region image out of prompt text and returns it as visual input', () => {    const annotation = pdfRegionAnnotation()    const sanitized = sanitizeAnnotations([annotation])     expect(sanitized).toEqual([annotation])    expect(annotationPayloadText(sanitized)).toContain('"type":"pdf-region"')    expect(annotationPayloadText(sanitized)).not.toContain('AQID')    expect(prepareAnnotationsForAgent('', sanitized).images).toEqual([      { mimeType: 'image/png', data: 'AQID', byteLength: 3 }    ])  })   it('sanitizes and serializes a Session item text annotation', () => {    const annotation = textAnnotation({      source: {        kind: 'session-item',        sessionId: '  session-1  ',        itemId: '  activity-1  ',        itemType: 'tool-activity',        sectionId: '  output  '      }    })     const sanitized = sanitizeAnnotations([annotation])     expect(sanitized).toEqual([      textAnnotation({        source: {          kind: 'session-item',          sessionId: 'session-1',          itemId: 'activity-1',          itemType: 'tool-activity',

Lines 1–500 of 603.