diff --git a/workspace-server/src/__tests__/services/DocsService.comments.test.ts b/workspace-server/src/__tests__/services/DocsService.comments.test.ts new file mode 100644 index 0000000..9ce09cc --- /dev/null +++ b/workspace-server/src/__tests__/services/DocsService.comments.test.ts @@ -0,0 +1,431 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + describe, + it, + expect, + jest, + beforeEach, + afterEach, +} from '@jest/globals'; +import { DocsService } from '../../services/DocsService'; +import { DriveService } from '../../services/DriveService'; +import { AuthManager } from '../../auth/AuthManager'; +import { google } from 'googleapis'; + +// Mock the googleapis module +jest.mock('googleapis'); +jest.mock('../../utils/logger'); +jest.mock('dompurify', () => { + return jest.fn().mockImplementation(() => ({ + sanitize: jest.fn((content) => content), + })); +}); + +describe('DocsService Comments and Suggestions', () => { + let docsService: DocsService; + let mockAuthManager: jest.Mocked; + let mockDriveService: jest.Mocked; + let mockDocsAPI: any; + let mockDriveAPI: any; + + beforeEach(() => { + jest.clearAllMocks(); + + mockAuthManager = { + getAuthenticatedClient: jest.fn(), + } as any; + + mockDriveService = { + findFolder: jest.fn(), + } as any; + + mockDocsAPI = { + documents: { + get: jest.fn(), + create: jest.fn(), + batchUpdate: jest.fn(), + }, + }; + + mockDriveAPI = { + files: { + create: jest.fn(), + list: jest.fn(), + get: jest.fn(), + update: jest.fn(), + }, + comments: { + list: jest.fn(), + }, + }; + + (google.docs as jest.Mock) = jest.fn().mockReturnValue(mockDocsAPI); + (google.drive as jest.Mock) = jest.fn().mockReturnValue(mockDriveAPI); + + docsService = new DocsService(mockAuthManager, mockDriveService); + + const mockAuthClient = { access_token: 'test-token' }; + mockAuthManager.getAuthenticatedClient.mockResolvedValue( + mockAuthClient as any, + ); + }); + + afterEach(() => { + jest.restoreAllMocks(); + }); + + describe('getSuggestions', () => { + it('should return suggestions as type text with JSON-stringified array', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + body: { + content: [ + { + paragraph: { + elements: [ + { + textRun: { + content: 'Suggested insertion', + suggestedInsertionIds: ['ins1'], + }, + startIndex: 1, + endIndex: 20, + }, + ], + }, + }, + ], + }, + }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + expect(result.content[0].type).toBe('text'); + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toHaveLength(1); + expect(suggestions[0]).toEqual({ + type: 'insertion', + text: 'Suggested insertion', + suggestionIds: ['ins1'], + startIndex: 1, + endIndex: 20, + }); + }); + + it('should extract insertion suggestions correctly', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + body: { + content: [ + { + paragraph: { + elements: [ + { + textRun: { + content: 'new text', + suggestedInsertionIds: ['sug-1', 'sug-2'], + }, + startIndex: 5, + endIndex: 13, + }, + ], + }, + }, + ], + }, + }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toHaveLength(1); + expect(suggestions[0].type).toBe('insertion'); + expect(suggestions[0].suggestionIds).toEqual(['sug-1', 'sug-2']); + }); + + it('should extract deletion suggestions correctly', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + body: { + content: [ + { + paragraph: { + elements: [ + { + textRun: { + content: 'deleted text', + suggestedDeletionIds: ['del-1'], + }, + startIndex: 1, + endIndex: 13, + }, + ], + }, + }, + ], + }, + }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toHaveLength(1); + expect(suggestions[0].type).toBe('deletion'); + expect(suggestions[0].text).toBe('deleted text'); + }); + + it('should extract style change suggestions correctly', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + body: { + content: [ + { + paragraph: { + elements: [ + { + textRun: { + content: 'styled text', + suggestedTextStyleChanges: { + 'style-1': { textStyle: { bold: true } }, + }, + textStyle: { bold: true }, + }, + startIndex: 1, + endIndex: 12, + }, + ], + }, + }, + ], + }, + }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toHaveLength(1); + expect(suggestions[0].type).toBe('styleChange'); + expect(suggestions[0].suggestionIds).toEqual(['style-1']); + expect(suggestions[0].textStyle).toEqual({ bold: true }); + }); + + it('should extract paragraph style change suggestions', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + body: { + content: [ + { + paragraph: { + paragraphStyle: { namedStyleType: 'HEADING_1' }, + suggestedParagraphStyleChanges: { + 'sug-para-1': { + paragraphStyle: { namedStyleType: 'HEADING_2' }, + }, + }, + elements: [ + { + textRun: { content: 'Heading Text' }, + startIndex: 1, + endIndex: 13, + }, + ], + }, + startIndex: 1, + endIndex: 13, + }, + ], + }, + }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toHaveLength(1); + expect(suggestions[0].type).toBe('paragraphStyleChange'); + expect(suggestions[0].suggestionIds).toEqual(['sug-para-1']); + expect(suggestions[0].namedStyleType).toBe('HEADING_2'); + expect(suggestions[0].text).toBe('Heading Text'); + }); + + it('should handle tables with recursive element processing', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + body: { + content: [ + { + table: { + tableRows: [ + { + tableCells: [ + { + content: [ + { + paragraph: { + elements: [ + { + textRun: { + content: 'cell text', + suggestedInsertionIds: ['cell-ins-1'], + }, + startIndex: 5, + endIndex: 14, + }, + ], + }, + }, + ], + }, + ], + }, + ], + }, + }, + ], + }, + }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toHaveLength(1); + expect(suggestions[0].type).toBe('insertion'); + expect(suggestions[0].text).toBe('cell text'); + }); + + it('should handle empty document body', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { body: null }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toEqual([]); + }); + + it('should handle API errors gracefully', async () => { + mockDocsAPI.documents.get.mockRejectedValue( + new Error('Docs API failed'), + ); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + expect(result.content[0].type).toBe('text'); + const parsed = JSON.parse(result.content[0].text); + expect(parsed).toEqual({ error: 'Docs API failed' }); + }); + + it('should handle undefined textRun.content with empty string fallback', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + body: { + content: [ + { + paragraph: { + elements: [ + { + textRun: { + content: undefined, + suggestedInsertionIds: ['ins-undef'], + }, + startIndex: 1, + endIndex: 5, + }, + ], + }, + }, + ], + }, + }, + }); + + const result = await docsService.getSuggestions({ + documentId: 'test-doc-id', + }); + + const suggestions = JSON.parse(result.content[0].text); + expect(suggestions).toHaveLength(1); + expect(suggestions[0].text).toBe(''); + }); + }); + + describe('getComments', () => { + it('should return comments as type text with JSON-stringified array', async () => { + const mockComments = [ + { + id: 'comment1', + content: 'This is a comment.', + author: { displayName: 'Test User', emailAddress: 'test@example.com' }, + createdTime: '2025-01-01T00:00:00Z', + resolved: false, + quotedFileContent: { value: 'quoted text' }, + }, + ]; + mockDriveAPI.comments.list.mockResolvedValue({ + data: { comments: mockComments }, + }); + + const result = await docsService.getComments({ + documentId: 'test-doc-id', + }); + + expect(result.content[0].type).toBe('text'); + const comments = JSON.parse(result.content[0].text); + expect(comments).toEqual(mockComments); + }); + + it('should handle empty comments list', async () => { + mockDriveAPI.comments.list.mockResolvedValue({ + data: { comments: [] }, + }); + + const result = await docsService.getComments({ + documentId: 'test-doc-id', + }); + + const comments = JSON.parse(result.content[0].text); + expect(comments).toEqual([]); + }); + + it('should handle API errors gracefully', async () => { + mockDriveAPI.comments.list.mockRejectedValue( + new Error('Comments API failed'), + ); + + const result = await docsService.getComments({ + documentId: 'test-doc-id', + }); + + expect(result.content[0].type).toBe('text'); + const parsed = JSON.parse(result.content[0].text); + expect(parsed).toEqual({ error: 'Comments API failed' }); + }); + }); +}); diff --git a/workspace-server/src/index.ts b/workspace-server/src/index.ts index becce91..f077b20 100644 --- a/workspace-server/src/index.ts +++ b/workspace-server/src/index.ts @@ -143,6 +143,28 @@ async function main() { }, ); + server.registerTool( + 'docs.getSuggestions', + { + description: 'Retrieves suggested edits from a Google Doc.', + inputSchema: { + documentId: z.string().describe('The ID of the document to retrieve suggestions from.'), + }, + }, + docsService.getSuggestions, + ); + + server.registerTool( + 'docs.getComments', + { + description: 'Retrieves comments from a Google Doc.', + inputSchema: { + documentId: z.string().describe('The ID of the document to retrieve comments from.'), + }, + }, + docsService.getComments, + ); + server.registerTool( 'docs.create', { diff --git a/workspace-server/src/services/DocsService.ts b/workspace-server/src/services/DocsService.ts index b62b7df..74eb97f 100644 --- a/workspace-server/src/services/DocsService.ts +++ b/workspace-server/src/services/DocsService.ts @@ -21,6 +21,16 @@ import { processMarkdownLineBreaks, } from '../utils/markdownToDocsRequests'; +interface DocsSuggestion { + type: 'insertion' | 'deletion' | 'styleChange' | 'paragraphStyleChange'; + text?: string; + suggestionIds?: string[]; + startIndex?: number; + endIndex?: number; + namedStyleType?: string; + textStyle?: docs_v1.Schema$TextStyle; +} + export class DocsService { private purify: ReturnType; @@ -44,6 +54,185 @@ export class DocsService { return google.drive({ version: 'v3', ...options }); } + public getSuggestions = async ({ + documentId, + }: { + documentId: string; + }) => { + logToFile( + `[DocsService] Starting getSuggestions for document: ${documentId}`, + ); + try { + const id = extractDocId(documentId) || documentId; + const docs = await this.getDocsClient(); + const res = await docs.documents.get({ + documentId: id, + suggestionsViewMode: 'SUGGESTIONS_INLINE', + fields: 'body', + }); + + const suggestions: DocsSuggestion[] = this._extractSuggestions(res.data.body); + + logToFile( + `[DocsService] Found ${suggestions.length} suggestions for document: ${id}`, + ); + + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify(suggestions, null, 2), + }, + ], + }; + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : String(error); + logToFile( + `[DocsService] Error during docs.getSuggestions: ${errorMessage}`, + ); + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify({ error: errorMessage }), + }, + ], + }; + } + }; + + private _extractSuggestions( + body: docs_v1.Schema$Body | undefined | null, + ): DocsSuggestion[] { + const suggestions: DocsSuggestion[] = []; + if (!body?.content) { + return suggestions; + } + + const processElements = ( + elements: docs_v1.Schema$StructuralElement[] | undefined, + ) => { + elements?.forEach((element) => { + if (element.paragraph) { + // Handle paragraph-level style suggestions + if (element.paragraph.suggestedParagraphStyleChanges) { + const suggestionIds = Object.keys( + element.paragraph.suggestedParagraphStyleChanges, + ); + if (suggestionIds.length > 0) { + const firstSuggestion = + element.paragraph.suggestedParagraphStyleChanges[ + suggestionIds[0] + ]; + suggestions.push({ + type: 'paragraphStyleChange', + text: this._getParagraphText(element.paragraph), + suggestionIds: suggestionIds, + namedStyleType: firstSuggestion?.paragraphStyle?.namedStyleType, + startIndex: element.startIndex, + endIndex: element.endIndex, + }); + } + } + + // Handle text-run-level suggestions within the paragraph + element.paragraph.elements?.forEach((pElement) => { + if (pElement.textRun) { + const baseSuggestion = { + text: pElement.textRun.content || '', + startIndex: pElement.startIndex, + endIndex: pElement.endIndex, + }; + + if (pElement.textRun.suggestedInsertionIds) { + suggestions.push({ + ...baseSuggestion, + type: 'insertion' as const, + suggestionIds: pElement.textRun.suggestedInsertionIds, + }); + } + if (pElement.textRun.suggestedDeletionIds) { + suggestions.push({ + ...baseSuggestion, + type: 'deletion' as const, + suggestionIds: pElement.textRun.suggestedDeletionIds, + }); + } + if (pElement.textRun.suggestedTextStyleChanges) { + suggestions.push({ + ...baseSuggestion, + type: 'styleChange' as const, + suggestionIds: Object.keys( + pElement.textRun.suggestedTextStyleChanges, + ), + textStyle: pElement.textRun.textStyle, + }); + } + } + }); + } else if (element.table) { + element.table.tableRows?.forEach((row) => { + row.tableCells?.forEach((cell) => { + processElements(cell.content); + }); + }); + } + }); + }; + + processElements(body.content); + return suggestions; + } + + private _getParagraphText(paragraph: docs_v1.Schema$Paragraph | undefined | null): string { + if (!paragraph?.elements) { + return ''; + } + return paragraph.elements + .map((pElement) => pElement.textRun?.content || '') + .join(''); + } + + public getComments = async ({ documentId }: { documentId: string }) => { + logToFile(`[DocsService] Starting getComments for document: ${documentId}`); + try { + const id = extractDocId(documentId) || documentId; + const drive = await this.getDriveClient(); + const res = await drive.comments.list({ + fileId: id, + fields: + 'comments(id, content, author(displayName, emailAddress), createdTime, resolved, quotedFileContent(value))', + }); + + const comments = res.data.comments || []; + logToFile( + `[DocsService] Found ${comments.length} comments for document: ${id}`, + ); + + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify(comments, null, 2), + }, + ], + }; + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : String(error); + logToFile(`[DocsService] Error during docs.getComments: ${errorMessage}`); + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify({ error: errorMessage }), + }, + ], + }; + } + }; + public create = async ({ title, folderName,