Agent: make AgentWorkspaceDocuments more robust

agent/src/AgentWorkspaceDocuments.test.ts

@@ -0,0 +1,84 @@
+import { beforeEach, describe, expect, it } from 'vitest'
+import * as vscode from 'vscode'
+import { ProtocolTextDocumentWithUri } from '../../vscode/src/jsonrpc/TextDocumentWithUri'
+import { AgentWorkspaceDocuments } from './AgentWorkspaceDocuments'
+
+describe('AgentWorkspaceDocuments', () => {
+    let documents: AgentWorkspaceDocuments
+    beforeEach(() => {
+        documents = new AgentWorkspaceDocuments()
+    })
+    const uri = vscode.Uri.parse('file:///foo.txt')
+
+    it('singleton document', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: 'hello' })
+        )
+        expect(document.getText()).toBe('hello')
+        const document2 = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: 'goodbye' })
+        )
+        // Regardless of when you got the reference to the document, `getText()`
+        // always reflects the latest value.
+        expect(document.getText()).toBe('goodbye')
+        expect(document2.getText()).toBe('goodbye')
+    })
+
+    it('null content', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: 'hello' })
+        )
+        expect(document.getText()).toBe('hello')
+        expect(documents.getDocument(uri)?.getText()).toBe('hello')
+
+        const document2 = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                contentChanges: null as any,
+                content: null as any,
+                visibleRange: null as any,
+                selection: null as any,
+            })
+        )
+        expect(document2.getText()).toBe('hello')
+        expect(document2.protocolDocument.contentChanges).toBeUndefined()
+        expect(document2.protocolDocument.selection).toBeUndefined()
+        expect(document2.protocolDocument.visibleRange).toBeUndefined()
+    })
+
+    it('incremental sync', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: ['abc', 'def', 'ghi'].join('\n') })
+        )
+        expect(document.getText()).toBe('abc\ndef\nghi')
+        documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                contentChanges: [
+                    {
+                        range: { start: { line: 0, character: 0 }, end: { line: 0, character: 1 } },
+                        text: 'x',
+                    },
+                    {
+                        range: { start: { line: 1, character: 1 }, end: { line: 1, character: 2 } },
+                        text: 'y',
+                    },
+                    {
+                        range: { start: { line: 2, character: 2 }, end: { line: 2, character: 3 } },
+                        text: 'z',
+                    },
+                ],
+            })
+        )
+        expect(document.getText()).toBe('xbc\ndyf\nghz')
+    })
+
+    it('visibleRanges', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                content: 'hello\ngoodbye\nworld\nsayonara\n',
+                visibleRange: { start: { line: 0, character: 0 }, end: { line: 1, character: 5 } },
+            })
+        )
+        const editor = documents.newTextEditor(document)
+        expect(editor.visibleRanges).toStrictEqual([new vscode.Selection(0, 0, 1, 5)])
+    })
+})

agent/src/AgentWorkspaceDocuments.ts

@@ -44,33 +44,53 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments {
     } {
         const fromCache = this.agentDocuments.get(document.underlying.uri)
         if (!fromCache) {
-            return { document: new AgentTextDocument(document), contentChanges: [] }
+            const result = new AgentTextDocument(document)
+            this.agentDocuments.set(document.underlying.uri, result)
+            return { document: result, contentChanges: [] }
         }
 
+        // We have seen this document before, which means we mutate the existing
+        // document to reflect the latest chagnes. For each URI, we keep a
+        // singleton document so that `AgentTextDocument.getText()` always
+        // reflects the latest value.
         const contentChanges: vscode.TextDocumentContentChangeEvent[] = []
-
         if (document.contentChanges && document.contentChanges.length > 0) {
+            // Incremental document sync.
             const changes = applyContentChanges(fromCache, document.contentChanges)
             contentChanges.push(...changes.contentChanges)
             document.underlying.content = changes.newText
-        } else if (document.content !== undefined && document.content != null) {
+        } else if (typeof document.content === 'string') {
+            // Full document sync.
             for (const change of calculateContentChanges(fromCache, document.content)) {
                 contentChanges.push(change)
             }
-        }
-
-        if (document.content === undefined) {
+        } else {
+            // No document sync. Use content from last update.
             document.underlying.content = fromCache.getText()
         }
 
-        if (document.selection === undefined) {
+        if (document.selection) {
             document.underlying.selection = fromCache.protocolDocument.selection
         }
 
-        if (document.visibleRange === undefined) {
+        if (document.visibleRange) {
             document.underlying.visibleRange = fromCache.protocolDocument.visibleRange
         }
 
+        // The client may send null values that we convert to undefined here.
+        if (document.content === null) {
+            document.underlying.content = undefined
+        }
+        if (document.contentChanges === null) {
+            document.underlying.contentChanges = undefined
+        }
+        if (document.selection === null) {
+            document.underlying.selection = undefined
+        }
+        if (document.visibleRange === null) {
+            document.underlying.visibleRange = undefined
+        }
+
         fromCache.update(document)
 
         return { document: fromCache, contentChanges }
@@ -107,7 +127,6 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments {
     } {
         const { document: agentDocument, contentChanges } =
             this.loadAndUpdateDocumentWithChanges(document)
-        this.agentDocuments.set(document.underlying.uri, agentDocument)
 
         const tabs: vscode.Tab[] = []
         for (const uri of this.allUris()) {