diff --git a/skills/google-docs/SKILL.md b/skills/google-docs/SKILL.md index eb0d1592..5c12afd1 100644 --- a/skills/google-docs/SKILL.md +++ b/skills/google-docs/SKILL.md @@ -189,8 +189,24 @@ docs.writeText({ ## Find and Replace -Use `docs.replaceText` to find all occurrences of a string and replace them. -This works across all tabs by default, or in a specific tab with `tabId`. +Use `docs.replaceText` only for literal replacements. It resolves actual Google +Docs indexes, edits from the highest index downward, and locks the current +revision so concurrent changes fail instead of shifting ranges. + +For section rewrites: + +1. Call `docs.previewReplaceSection` with a tab ID and unique heading anchors. +2. Review its `currentText`, `range`, and `revisionId`. +3. Call `docs.replaceSection` with that exact `expectedText` and + `expectedRevisionId`, plus replacement paragraph blocks. + +The write is rejected if headings are ambiguous or the document changed after +preview. Never substitute `docs.writeText` plus manually calculated indexes for +an existing rich document. + +Use `docs.getStructure` when exact paragraph ranges, heading styles, or the +current revision are needed. Unlike `docs.getText`, it returns API indexes +instead of rendered-text offsets. ## Tab Management diff --git a/workspace-server/src/__tests__/features/feature-resolver.test.ts b/workspace-server/src/__tests__/features/feature-resolver.test.ts index fa473f3c..defdf7d7 100644 --- a/workspace-server/src/__tests__/features/feature-resolver.test.ts +++ b/workspace-server/src/__tests__/features/feature-resolver.test.ts @@ -85,6 +85,14 @@ describe('resolveFeatures', () => { } }); + it('enables safe structural Docs tools by default', () => { + const { enabledTools } = resolveFeatures(); + + expect(enabledTools.has('docs.getStructure')).toBe(true); + expect(enabledTools.has('docs.previewReplaceSection')).toBe(true); + expect(enabledTools.has('docs.replaceSection')).toBe(true); + }); + it('should compute scopes only for enabled groups', () => { const { requiredScopes } = resolveFeatures(); diff --git a/workspace-server/src/__tests__/integration/DocsService.integration.test.ts b/workspace-server/src/__tests__/integration/DocsService.integration.test.ts new file mode 100644 index 00000000..6110d448 --- /dev/null +++ b/workspace-server/src/__tests__/integration/DocsService.integration.test.ts @@ -0,0 +1,171 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { afterAll, beforeAll, describe, expect, it } from '@jest/globals'; +import { google } from 'googleapis'; +import { AuthManager } from '../../auth/AuthManager'; +import { SCOPES } from '../../auth/scopes'; +import { DocsService } from '../../services/DocsService'; + +const describeIntegration = + process.env['GOOGLE_DOCS_INTEGRATION'] === '1' ? describe : describe.skip; + +describeIntegration('DocsService live integration', () => { + const authManager = new AuthManager(SCOPES); + const docsService = new DocsService(authManager); + let documentId: string; + let primaryTabId: string; + + beforeAll(async () => { + const auth = await authManager.getAuthenticatedClient(); + const docs = google.docs({ version: 'v1', auth }); + const created = await docs.documents.create({ + requestBody: { + title: `Safe Docs MCP integration ${new Date().toISOString()}`, + }, + }); + documentId = created.data.documentId!; + + const initial = await docs.documents.get({ + documentId, + includeTabsContent: true, + fields: 'tabs(tabProperties)', + }); + primaryTabId = initial.data.tabs?.[0]?.tabProperties?.tabId || ''; + if (!primaryTabId) { + throw new Error('Fixture document did not expose its primary tab ID.'); + } + + await docs.documents.batchUpdate({ + documentId, + requestBody: { + requests: [ + { + insertText: { + location: { tabId: primaryTabId, index: 1 }, + text: 'Fixture\nConcepts\nOld body\nGrowth scenarios\nOutside sentinel\n', + }, + }, + { + updateParagraphStyle: { + range: { tabId: primaryTabId, startIndex: 9, endIndex: 18 }, + paragraphStyle: { namedStyleType: 'HEADING_2' }, + fields: 'namedStyleType', + }, + }, + { + updateParagraphStyle: { + range: { tabId: primaryTabId, startIndex: 27, endIndex: 44 }, + paragraphStyle: { namedStyleType: 'HEADING_2' }, + fields: 'namedStyleType', + }, + }, + { + insertRichLink: { + location: { tabId: primaryTabId, index: 1 }, + richLinkProperties: { + uri: `https://docs.google.com/document/d/${documentId}/edit`, + }, + }, + }, + { + insertTable: { + rows: 1, + columns: 1, + endOfSegmentLocation: { tabId: primaryTabId }, + }, + }, + { + addDocumentTab: { + tabProperties: { title: 'Secondary fixture tab' }, + }, + }, + ], + }, + }); + + const withTabs = await docs.documents.get({ + documentId, + includeTabsContent: true, + fields: 'tabs(tabProperties)', + }); + const secondaryTabId = withTabs.data.tabs?.find( + (tab) => tab.tabProperties?.title === 'Secondary fixture tab', + )?.tabProperties?.tabId; + if (!secondaryTabId) { + throw new Error('Fixture document did not create its secondary tab.'); + } + await docs.documents.batchUpdate({ + documentId, + requestBody: { + requests: [ + { + insertText: { + location: { tabId: secondaryTabId, index: 1 }, + text: 'Secondary sentinel\n', + }, + }, + ], + }, + }); + }, 120_000); + + afterAll(async () => { + if (!documentId) { + return; + } + const auth = await authManager.getAuthenticatedClient(); + const drive = google.drive({ version: 'v3', auth }); + await drive.files.update({ + fileId: documentId, + requestBody: { trashed: true }, + supportsAllDrives: true, + }); + }, 120_000); + + it('replaces only the selected section in a rich multi-tab document', async () => { + const preview = await docsService.previewReplaceSection({ + documentId, + tabId: primaryTabId, + startHeading: 'Concepts', + endHeading: 'Growth scenarios', + }); + const previewBody = JSON.parse(preview.content[0].text); + expect(previewBody.error).toBeUndefined(); + + const replacement = await docsService.replaceSection({ + documentId, + tabId: primaryTabId, + startHeading: 'Concepts', + endHeading: 'Growth scenarios', + expectedRevisionId: previewBody.revisionId, + expectedText: previewBody.currentText, + blocks: [{ text: 'Concepts', style: 'heading2' }, { text: 'New body' }], + }); + expect(JSON.parse(replacement.content[0].text).replaced).toBe(true); + + const auth = await authManager.getAuthenticatedClient(); + const docs = google.docs({ version: 'v1', auth }); + const current = await docs.documents.get({ + documentId, + includeTabsContent: true, + fields: 'tabs(tabProperties,documentTab(body))', + }); + const primary = current.data.tabs?.find( + (tab) => tab.tabProperties?.tabId === primaryTabId, + ); + const primaryText = ( + await docsService.getText({ documentId, tabId: primaryTabId }) + ).content[0].text; + + expect(primaryText).toContain('Concepts\nNew body\n'); + expect(primaryText).toContain('Growth scenarios\nOutside sentinel\n'); + expect( + primary?.documentTab?.body?.content?.some((element) => element.table), + ).toBe(true); + expect(current.data.tabs).toHaveLength(2); + }, 120_000); +}); diff --git a/workspace-server/src/__tests__/services/DocsService.test.ts b/workspace-server/src/__tests__/services/DocsService.test.ts index d3aa66e6..72db9b8f 100644 --- a/workspace-server/src/__tests__/services/DocsService.test.ts +++ b/workspace-server/src/__tests__/services/DocsService.test.ts @@ -12,7 +12,11 @@ import { beforeEach, afterEach, } from '@jest/globals'; -import { DocsService, TABS_FIELD_MASK } from '../../services/DocsService'; +import { + DocsService, + SAFE_DOCS_FIELD_MASK, + TABS_FIELD_MASK, +} from '../../services/DocsService'; import { AuthManager } from '../../auth/AuthManager'; import { google } from 'googleapis'; @@ -816,6 +820,7 @@ describe('DocsService', () => { // Mock the document get call that finds occurrences mockDocsAPI.documents.get.mockResolvedValue({ data: { + revisionId: 'rev-1', tabs: [ { documentTab: { @@ -851,7 +856,7 @@ describe('DocsService', () => { expect(mockDocsAPI.documents.get).toHaveBeenCalledWith({ documentId: 'test-doc-id', - fields: TABS_FIELD_MASK, + fields: SAFE_DOCS_FIELD_MASK, includeTabsContent: true, }); @@ -878,17 +883,21 @@ describe('DocsService', () => { }, }), ]), + writeControl: { requiredRevisionId: 'rev-1' }, }, }); - expect(result.content[0].text).toBe( - 'Successfully replaced text in document test-doc-id', - ); + expect(JSON.parse(result.content[0].text)).toMatchObject({ + documentId: 'test-doc-id', + revisionId: 'rev-1', + matchCount: 2, + }); }); it('should replace text with literal content (no markdown parsing)', async () => { // Mock the document get call that finds occurrences mockDocsAPI.documents.get.mockResolvedValue({ data: { + revisionId: 'rev-1', tabs: [ { documentTab: { @@ -928,7 +937,7 @@ describe('DocsService', () => { expect(mockDocsAPI.documents.get).toHaveBeenCalledWith({ documentId: 'test-doc-id', - fields: TABS_FIELD_MASK, + fields: SAFE_DOCS_FIELD_MASK, includeTabsContent: true, }); @@ -937,13 +946,13 @@ describe('DocsService', () => { documentId: 'test-doc-id', requestBody: { requests: [ - // First occurrence + // Highest index first, so earlier ranges do not move. { deleteContentRange: { range: { tabId: undefined, - startIndex: 9, - endIndex: 18, + startIndex: 23, + endIndex: 32, }, }, }, @@ -951,18 +960,17 @@ describe('DocsService', () => { insertText: { location: { tabId: undefined, - index: 9, + index: 23, }, text: '**bold text**', }, }, - // Second occurrence (offset by length diff: 13 - 9 = +4) { deleteContentRange: { range: { tabId: undefined, - startIndex: 27, - endIndex: 36, + startIndex: 9, + endIndex: 18, }, }, }, @@ -970,23 +978,23 @@ describe('DocsService', () => { insertText: { location: { tabId: undefined, - index: 27, + index: 9, }, text: '**bold text**', }, }, ], + writeControl: { requiredRevisionId: 'rev-1' }, }, }); - expect(result.content[0].text).toBe( - 'Successfully replaced text in document test-doc-id', - ); + expect(JSON.parse(result.content[0].text).matchCount).toBe(2); }); it('should handle errors during replaceText', async () => { // Mock the document get call mockDocsAPI.documents.get.mockResolvedValue({ data: { + revisionId: 'rev-1', tabs: [ { documentTab: { @@ -1022,6 +1030,7 @@ describe('DocsService', () => { it('should replace text in a specific tab using delete/insert', async () => { mockDocsAPI.documents.get.mockResolvedValue({ data: { + revisionId: 'rev-1', tabs: [ { tabProperties: { tabId: 'tab-1' }, @@ -1076,6 +1085,7 @@ describe('DocsService', () => { }, }), ]), + writeControl: { requiredRevisionId: 'rev-1' }, }, }); }); @@ -1083,6 +1093,7 @@ describe('DocsService', () => { it('should replace text in a nested child tab by tabId', async () => { mockDocsAPI.documents.get.mockResolvedValue({ data: { + revisionId: 'rev-1', tabs: [ { tabProperties: { tabId: 'parent-tab' }, @@ -1154,8 +1165,533 @@ describe('DocsService', () => { }, }), ]), + writeControl: { requiredRevisionId: 'rev-1' }, + }, + }); + }); + + it('uses API indexes after a rich link and locks the document revision', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + revisionId: 'rev-1', + tabs: [ + { + tabProperties: { tabId: 'tab-1' }, + documentTab: { + body: { + content: [ + { + startIndex: 1, + endIndex: 25, + paragraph: { + elements: [ + { + startIndex: 1, + endIndex: 2, + richLink: { + richLinkProperties: { + title: 'Link', + uri: 'https://example.com', + }, + }, + }, + { + startIndex: 10, + endIndex: 25, + textRun: { content: 'Replace target\n' }, + }, + ], + }, + }, + ], + }, + }, + }, + ], }, }); + mockDocsAPI.documents.batchUpdate.mockResolvedValue({ data: {} }); + + const result = await docsService.replaceText({ + documentId: 'test-doc-id', + findText: 'target', + replaceText: 'value', + tabId: 'tab-1', + }); + + expect(mockDocsAPI.documents.batchUpdate).toHaveBeenCalledWith({ + documentId: 'test-doc-id', + requestBody: { + requests: [ + { + deleteContentRange: { + range: { + tabId: 'tab-1', + startIndex: 18, + endIndex: 24, + }, + }, + }, + { + insertText: { + location: { tabId: 'tab-1', index: 18 }, + text: 'value', + }, + }, + ], + writeControl: { requiredRevisionId: 'rev-1' }, + }, + }); + expect(JSON.parse(result.content[0].text)).toMatchObject({ + matchCount: 1, + revisionId: 'rev-1', + }); + }); + + it('matches text split across contiguous text runs', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + revisionId: 'rev-1', + tabs: [ + { + tabProperties: { tabId: 'tab-1' }, + documentTab: { + body: { + content: [ + { + startIndex: 100, + endIndex: 107, + paragraph: { + elements: [ + { + startIndex: 100, + endIndex: 103, + textRun: { content: 'tar' }, + }, + { + startIndex: 103, + endIndex: 107, + textRun: { content: 'get\n' }, + }, + ], + }, + }, + ], + }, + }, + }, + ], + }, + }); + mockDocsAPI.documents.batchUpdate.mockResolvedValue({ data: {} }); + + const result = await docsService.replaceText({ + documentId: 'test-doc-id', + findText: 'target', + replaceText: 'value', + tabId: 'tab-1', + }); + + expect(mockDocsAPI.documents.batchUpdate).toHaveBeenCalledWith({ + documentId: 'test-doc-id', + requestBody: { + requests: [ + { + deleteContentRange: { + range: { + tabId: 'tab-1', + startIndex: 100, + endIndex: 106, + }, + }, + }, + { + insertText: { + location: { tabId: 'tab-1', index: 100 }, + text: 'value', + }, + }, + ], + writeControl: { requiredRevisionId: 'rev-1' }, + }, + }); + expect(JSON.parse(result.content[0].text).matchCount).toBe(1); + }); + + it('uses actual indexes for text inside a table', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + revisionId: 'rev-1', + tabs: [ + { + tabProperties: { tabId: 'tab-1' }, + documentTab: { + body: { + content: [ + { + startIndex: 1, + endIndex: 80, + table: { + tableRows: [ + { + tableCells: [ + { + content: [ + { + startIndex: 50, + endIndex: 57, + paragraph: { + elements: [ + { + startIndex: 50, + endIndex: 57, + textRun: { content: 'target\n' }, + }, + ], + }, + }, + ], + }, + ], + }, + ], + }, + }, + ], + }, + }, + }, + ], + }, + }); + mockDocsAPI.documents.batchUpdate.mockResolvedValue({ data: {} }); + + await docsService.replaceText({ + documentId: 'test-doc-id', + findText: 'target', + replaceText: 'value', + tabId: 'tab-1', + }); + + expect(mockDocsAPI.documents.batchUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + requestBody: expect.objectContaining({ + requests: expect.arrayContaining([ + { + deleteContentRange: { + range: { + tabId: 'tab-1', + startIndex: 50, + endIndex: 56, + }, + }, + }, + ]), + }), + }), + ); + }); + + it('returns an explicit no-op when text is absent', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + revisionId: 'rev-1', + tabs: [ + { + tabProperties: { tabId: 'tab-1' }, + documentTab: { + body: { + content: [ + { + startIndex: 1, + endIndex: 6, + paragraph: { + elements: [ + { + startIndex: 1, + endIndex: 6, + textRun: { content: 'other' }, + }, + ], + }, + }, + ], + }, + }, + }, + ], + }, + }); + + const result = await docsService.replaceText({ + documentId: 'test-doc-id', + findText: 'target', + replaceText: 'value', + tabId: 'tab-1', + }); + + expect(JSON.parse(result.content[0].text)).toMatchObject({ + revisionId: 'rev-1', + matchCount: 0, + }); + expect(mockDocsAPI.documents.batchUpdate).not.toHaveBeenCalled(); + }); + }); + + describe('safe structural editing', () => { + const structuredDocument = { + data: { + title: 'Quota RFC', + revisionId: 'rev-1', + tabs: [ + { + tabProperties: { tabId: 'tab-1', title: 'Quota spec' }, + documentTab: { + body: { + content: [ + { + startIndex: 20, + endIndex: 29, + paragraph: { + paragraphStyle: { namedStyleType: 'HEADING_2' }, + elements: [ + { + startIndex: 20, + endIndex: 29, + textRun: { content: 'Concepts\n' }, + }, + ], + }, + }, + { + startIndex: 29, + endIndex: 33, + paragraph: { + paragraphStyle: { namedStyleType: 'NORMAL_TEXT' }, + elements: [ + { + startIndex: 29, + endIndex: 33, + textRun: { content: 'Old\n' }, + }, + ], + }, + }, + { + startIndex: 33, + endIndex: 50, + paragraph: { + paragraphStyle: { namedStyleType: 'HEADING_2' }, + elements: [ + { + startIndex: 33, + endIndex: 50, + textRun: { content: 'Growth scenarios\n' }, + }, + ], + }, + }, + ], + }, + }, + }, + ], + }, + }; + + it('returns tab-aware paragraph ranges and revision ID', async () => { + mockDocsAPI.documents.get.mockResolvedValue(structuredDocument); + + const result = await docsService.getStructure({ + documentId: 'test-doc-id', + tabId: 'tab-1', + }); + + expect(JSON.parse(result.content[0].text)).toEqual({ + documentId: 'test-doc-id', + title: 'Quota RFC', + revisionId: 'rev-1', + tabs: [ + { + tabId: 'tab-1', + title: 'Quota spec', + paragraphs: [ + { + startIndex: 20, + endIndex: 29, + style: 'HEADING_2', + text: 'Concepts\n', + }, + { + startIndex: 29, + endIndex: 33, + style: 'NORMAL_TEXT', + text: 'Old\n', + }, + { + startIndex: 33, + endIndex: 50, + style: 'HEADING_2', + text: 'Growth scenarios\n', + }, + ], + }, + ], + }); + }); + + it('previews a uniquely anchored section without writing', async () => { + mockDocsAPI.documents.get.mockResolvedValue(structuredDocument); + + const result = await docsService.previewReplaceSection({ + documentId: 'test-doc-id', + tabId: 'tab-1', + startHeading: 'Concepts', + endHeading: 'Growth scenarios', + }); + + expect(JSON.parse(result.content[0].text)).toMatchObject({ + revisionId: 'rev-1', + range: { tabId: 'tab-1', startIndex: 20, endIndex: 33 }, + currentText: 'Concepts\nOld\n', + }); + expect(mockDocsAPI.documents.batchUpdate).not.toHaveBeenCalled(); + }); + + it('replaces a previewed section atomically with revision locking', async () => { + mockDocsAPI.documents.get.mockResolvedValue(structuredDocument); + mockDocsAPI.documents.batchUpdate.mockResolvedValue({ data: {} }); + + const result = await docsService.replaceSection({ + documentId: 'test-doc-id', + tabId: 'tab-1', + startHeading: 'Concepts', + endHeading: 'Growth scenarios', + expectedRevisionId: 'rev-1', + expectedText: 'Concepts\nOld\n', + blocks: [ + { text: 'Concepts', style: 'HEADING_2' }, + { text: 'New body' }, + ], + }); + + expect(mockDocsAPI.documents.batchUpdate).toHaveBeenCalledWith({ + documentId: 'test-doc-id', + requestBody: { + requests: [ + { + deleteContentRange: { + range: { + tabId: 'tab-1', + startIndex: 20, + endIndex: 33, + }, + }, + }, + { + insertText: { + location: { tabId: 'tab-1', index: 20 }, + text: 'Concepts\nNew body\n', + }, + }, + { + updateParagraphStyle: { + range: { + tabId: 'tab-1', + startIndex: 20, + endIndex: 29, + }, + paragraphStyle: { namedStyleType: 'HEADING_2' }, + fields: 'namedStyleType', + }, + }, + ], + writeControl: { requiredRevisionId: 'rev-1' }, + }, + }); + expect(JSON.parse(result.content[0].text)).toMatchObject({ + replaced: true, + revisionId: 'rev-1', + range: { tabId: 'tab-1', startIndex: 20, endIndex: 33 }, + }); + }); + + it('refuses replacement when the preview revision is stale', async () => { + mockDocsAPI.documents.get.mockResolvedValue({ + ...structuredDocument, + data: { ...structuredDocument.data, revisionId: 'rev-2' }, + }); + + const result = await docsService.replaceSection({ + documentId: 'test-doc-id', + tabId: 'tab-1', + startHeading: 'Concepts', + endHeading: 'Growth scenarios', + expectedRevisionId: 'rev-1', + expectedText: 'Concepts\nOld\n', + blocks: [{ text: 'Concepts', style: 'HEADING_2' }], + }); + + expect(JSON.parse(result.content[0].text).error).toContain( + 'revision changed', + ); + expect(mockDocsAPI.documents.batchUpdate).not.toHaveBeenCalled(); + }); + + it('refuses replacement when the section text differs from the preview', async () => { + mockDocsAPI.documents.get.mockResolvedValue(structuredDocument); + + const result = await docsService.replaceSection({ + documentId: 'test-doc-id', + tabId: 'tab-1', + startHeading: 'Concepts', + endHeading: 'Growth scenarios', + expectedRevisionId: 'rev-1', + expectedText: 'Concepts\nDifferent body\n', + blocks: [{ text: 'Concepts', style: 'HEADING_2' }], + }); + + expect(JSON.parse(result.content[0].text).error).toContain( + 'section text changed', + ); + expect(mockDocsAPI.documents.batchUpdate).not.toHaveBeenCalled(); + }); + + it('refuses ambiguous section headings', async () => { + const duplicateHeading = + structuredDocument.data.tabs[0].documentTab.body.content[0]; + mockDocsAPI.documents.get.mockResolvedValue({ + data: { + ...structuredDocument.data, + tabs: [ + { + ...structuredDocument.data.tabs[0], + documentTab: { + body: { + content: [ + ...structuredDocument.data.tabs[0].documentTab.body.content, + { + ...duplicateHeading, + startIndex: 60, + endIndex: 69, + }, + ], + }, + }, + }, + ], + }, + }); + + const result = await docsService.previewReplaceSection({ + documentId: 'test-doc-id', + tabId: 'tab-1', + startHeading: 'Concepts', + endHeading: 'Growth scenarios', + }); + + expect(JSON.parse(result.content[0].text).error).toContain('found 2'); + expect(mockDocsAPI.documents.batchUpdate).not.toHaveBeenCalled(); }); }); }); diff --git a/workspace-server/src/features/feature-config.ts b/workspace-server/src/features/feature-config.ts index dc9eef53..cccd9d7d 100644 --- a/workspace-server/src/features/feature-config.ts +++ b/workspace-server/src/features/feature-config.ts @@ -61,7 +61,12 @@ export const FEATURE_GROUPS: readonly FeatureGroup[] = [ service: 'docs', group: 'read', scopes: scopes('documents'), - tools: ['docs.getSuggestions', 'docs.getText'], + tools: [ + 'docs.getSuggestions', + 'docs.getText', + 'docs.getStructure', + 'docs.previewReplaceSection', + ], defaultEnabled: true, }, { @@ -72,6 +77,7 @@ export const FEATURE_GROUPS: readonly FeatureGroup[] = [ 'docs.create', 'docs.writeText', 'docs.replaceText', + 'docs.replaceSection', 'docs.formatText', ], defaultEnabled: true, diff --git a/workspace-server/src/index.ts b/workspace-server/src/index.ts index e7a8077b..fce10873 100644 --- a/workspace-server/src/index.ts +++ b/workspace-server/src/index.ts @@ -379,11 +379,115 @@ async function main() { docsService.getText, ); + registerTool( + 'docs.getStructure', + { + description: + 'Retrieves tab-aware paragraph text, styles, exact Google Docs indexes, and the current revision ID for safe editing.', + inputSchema: { + documentId: z.string().describe('The ID or URL of the document.'), + tabId: z + .string() + .optional() + .describe('The tab to inspect. If omitted, returns every tab.'), + }, + ...readOnlyToolProps, + }, + docsService.getStructure, + ); + + registerTool( + 'docs.previewReplaceSection', + { + description: + 'Resolves a uniquely headed section and returns its exact range, text, and revision without writing.', + inputSchema: { + documentId: z.string().describe('The ID or URL of the document.'), + tabId: z.string().describe('The tab containing the section.'), + startHeading: z + .string() + .trim() + .min(1) + .describe('Exact heading text at the start of the replacement.'), + endHeading: z + .string() + .trim() + .min(1) + .optional() + .describe( + 'Exact heading after the section. If omitted, replaces through the end of the tab.', + ), + }, + ...readOnlyToolProps, + }, + docsService.previewReplaceSection, + ); + + registerTool( + 'docs.replaceSection', + { + description: + 'Atomically replaces one uniquely headed section using the revision returned by previewReplaceSection.', + inputSchema: { + documentId: z.string().describe('The ID or URL of the document.'), + tabId: z.string().describe('The tab containing the section.'), + startHeading: z + .string() + .trim() + .min(1) + .describe('Exact heading text at the start of the replacement.'), + endHeading: z + .string() + .trim() + .min(1) + .optional() + .describe( + 'Exact heading after the section. If omitted, replaces through the end of the tab.', + ), + expectedRevisionId: z + .string() + .min(1) + .describe('Revision ID returned by the preview operation.'), + expectedText: z + .string() + .describe('Exact currentText returned by the preview operation.'), + blocks: z + .array( + z.object({ + text: z.string().describe('Paragraph text.'), + style: z + .enum([ + 'heading1', + 'heading2', + 'heading3', + 'heading4', + 'heading5', + 'heading6', + 'normalText', + 'HEADING_1', + 'HEADING_2', + 'HEADING_3', + 'HEADING_4', + 'HEADING_5', + 'HEADING_6', + 'NORMAL_TEXT', + ]) + .optional() + .describe('Optional Google Docs paragraph style.'), + }), + ) + .min(1) + .describe('Replacement paragraphs in document order.'), + }, + }, + docsService.replaceSection, + ); + registerTool( 'docs.replaceText', { description: - 'Replaces all occurrences of a given text with new text in a Google Doc.', + 'Safely replaces text using real Docs indexes and revision locking. Returns the number of matches.', inputSchema: { documentId: z.string().describe('The ID of the document to modify.'), findText: z.string().describe('The text to find in the document.'), diff --git a/workspace-server/src/services/DocsService.ts b/workspace-server/src/services/DocsService.ts index 72350bb7..f3f00a6c 100644 --- a/workspace-server/src/services/DocsService.ts +++ b/workspace-server/src/services/DocsService.ts @@ -16,6 +16,7 @@ import { extractDocumentId as validateAndExtractDocId } from '../utils/validatio // "comment-specific fields" errors when combined with includeTabsContent. export const TABS_FIELD_MASK = 'tabs(tabProperties,documentTab(body,headers,footers,footnotes))'; +export const SAFE_DOCS_FIELD_MASK = `title,revisionId,${TABS_FIELD_MASK}`; interface BaseDocsSuggestion { text: string; @@ -51,6 +52,27 @@ type DocsSuggestion = | DocsStyleChangeSuggestion | DocsParagraphStyleChangeSuggestion; +interface DocsReplacementBlock { + text: string; + style?: string; +} + +interface StructuredParagraph { + startIndex: number; + endIndex: number; + style: string; + text: string; +} + +interface SectionLocation { + range: { + tabId: string; + startIndex: number; + endIndex: number; + }; + currentText: string; +} + export class DocsService { /** * Recursively flattens a tab tree into a single array, @@ -722,6 +744,173 @@ export class DocsService { return displayText || timestamp || ''; } + public getStructure = async ({ + documentId, + tabId, + }: { + documentId: string; + tabId?: string; + }) => { + try { + const id = validateAndExtractDocId(documentId); + const docs = await this.getDocsClient(); + const response = await docs.documents.get({ + documentId: id, + fields: SAFE_DOCS_FIELD_MASK, + includeTabsContent: true, + suggestionsViewMode: 'PREVIEW_WITHOUT_SUGGESTIONS', + }); + const revisionId = response.data.revisionId; + if (!revisionId) { + throw new Error('Google Docs response did not include a revision ID.'); + } + + let tabs = this._flattenTabs(response.data.tabs || []); + if (tabId) { + tabs = tabs.filter((tab) => tab.tabProperties?.tabId === tabId); + if (tabs.length === 0) { + throw new Error(`Tab with ID ${tabId} not found.`); + } + } + + const result = { + documentId: id, + title: response.data.title, + revisionId, + tabs: tabs.map((tab) => ({ + tabId: tab.tabProperties?.tabId, + title: tab.tabProperties?.title, + paragraphs: this._collectParagraphs(tab.documentTab?.body?.content), + })), + }; + return { + content: [{ type: 'text' as const, text: JSON.stringify(result) }], + }; + } catch (error) { + return this._docsError('docs.getStructure', error); + } + }; + + public previewReplaceSection = async ({ + documentId, + tabId, + startHeading, + endHeading, + }: { + documentId: string; + tabId: string; + startHeading: string; + endHeading?: string; + }) => { + try { + const loaded = await this._loadStructuredDocument(documentId); + const section = this._resolveSection( + loaded.tabs, + tabId, + startHeading, + endHeading, + ); + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify({ + documentId: loaded.id, + revisionId: loaded.revisionId, + ...section, + }), + }, + ], + }; + } catch (error) { + return this._docsError('docs.previewReplaceSection', error); + } + }; + + public replaceSection = async ({ + documentId, + tabId, + startHeading, + endHeading, + expectedRevisionId, + expectedText, + blocks, + }: { + documentId: string; + tabId: string; + startHeading: string; + endHeading?: string; + expectedRevisionId: string; + expectedText: string; + blocks: DocsReplacementBlock[]; + }) => { + try { + if (blocks.length === 0) { + throw new Error('At least one replacement block is required.'); + } + const loaded = await this._loadStructuredDocument(documentId); + if (loaded.revisionId !== expectedRevisionId) { + throw new Error( + `Document revision changed from ${expectedRevisionId} to ${loaded.revisionId}; preview the section again.`, + ); + } + const section = this._resolveSection( + loaded.tabs, + tabId, + startHeading, + endHeading, + ); + if (section.currentText !== expectedText) { + throw new Error( + 'The section text changed after preview; preview the section again.', + ); + } + const rendered = this._renderBlocks( + blocks, + section.range.startIndex, + tabId, + ); + const requests: docs_v1.Schema$Request[] = [ + { deleteContentRange: { range: section.range } }, + { + insertText: { + location: { + tabId, + index: section.range.startIndex, + }, + text: rendered.text, + }, + }, + ...rendered.styleRequests.map((styleRequest) => ({ + updateParagraphStyle: styleRequest, + })), + ]; + + await loaded.docs.documents.batchUpdate({ + documentId: loaded.id, + requestBody: { + requests, + writeControl: { requiredRevisionId: expectedRevisionId }, + }, + }); + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify({ + replaced: true, + documentId: loaded.id, + revisionId: expectedRevisionId, + range: section.range, + }), + }, + ], + }; + } catch (error) { + return this._docsError('docs.replaceSection', error); + } + }; + public replaceText = async ({ documentId, findText, @@ -743,9 +932,13 @@ export class DocsService { // Get the document to find where the text will be replaced const docBefore = await docs.documents.get({ documentId: id, - fields: TABS_FIELD_MASK, + fields: SAFE_DOCS_FIELD_MASK, includeTabsContent: true, }); + const revisionId = docBefore.data.revisionId; + if (!revisionId) { + throw new Error('Google Docs response did not include a revision ID.'); + } const tabs = this._flattenTabs(docBefore.data.tabs || []); @@ -785,6 +978,7 @@ export class DocsService { documentId: id, requestBody: { requests, + writeControl: { requiredRevisionId: revisionId }, }, }); } @@ -794,7 +988,11 @@ export class DocsService { content: [ { type: 'text' as const, - text: `Successfully replaced text in document ${id}`, + text: JSON.stringify({ + documentId: id, + revisionId, + matchCount: requests.length / 2, + }), }, ], }; @@ -819,58 +1017,242 @@ export class DocsService { findText: string, newText: string, ): docs_v1.Schema$Request[] { - const requests: docs_v1.Schema$Request[] = []; - const documentText = this._getFullDocumentText(content); - const occurrences: number[] = []; - let searchIndex = 0; - while ((searchIndex = documentText.indexOf(findText, searchIndex)) !== -1) { - occurrences.push(searchIndex + 1); - searchIndex += findText.length; + if (!findText) { + throw new Error('findText must not be empty.'); } - - const lengthDiff = newText.length - findText.length; - let cumulativeOffset = 0; - - for (let i = 0; i < occurrences.length; i++) { - const occurrence = occurrences[i]; - const adjustedPosition = occurrence + cumulativeOffset; - - // Delete old text + const matches = this._findTextMatches(content, findText).sort( + (left, right) => right.startIndex - left.startIndex, + ); + const requests: docs_v1.Schema$Request[] = []; + for (const match of matches) { requests.push({ deleteContentRange: { range: { tabId: tabId, - startIndex: adjustedPosition, - endIndex: adjustedPosition + findText.length, + startIndex: match.startIndex, + endIndex: match.endIndex, }, }, }); - - // Insert new text requests.push({ insertText: { location: { tabId: tabId, - index: adjustedPosition, + index: match.startIndex, }, text: newText, }, }); - - cumulativeOffset += lengthDiff; } return requests; } - private _getFullDocumentText( + private _findTextMatches( content: docs_v1.Schema$StructuralElement[] | undefined, - ): string { + findText: string, + ): { startIndex: number; endIndex: number }[] { + const matches: { startIndex: number; endIndex: number }[] = []; + for (const element of content || []) { + if (element.paragraph) { + const chunks: { startIndex: number; endIndex: number; text: string }[] = + []; + let fallbackIndex = element.startIndex ?? 1; + for (const paragraphElement of element.paragraph.elements || []) { + const runText = paragraphElement.textRun?.content; + if (runText === undefined) { + fallbackIndex = paragraphElement.endIndex ?? fallbackIndex; + continue; + } + const startIndex = paragraphElement.startIndex ?? fallbackIndex; + const endIndex = + paragraphElement.endIndex ?? startIndex + runText.length; + const previous = chunks[chunks.length - 1]; + if (previous && previous.endIndex === startIndex) { + previous.text += runText; + previous.endIndex = endIndex; + } else { + chunks.push({ startIndex, endIndex, text: runText }); + } + fallbackIndex = endIndex; + } + for (const chunk of chunks) { + let offset = 0; + while ((offset = chunk.text.indexOf(findText, offset)) !== -1) { + matches.push({ + startIndex: chunk.startIndex + offset, + endIndex: chunk.startIndex + offset + findText.length, + }); + offset += findText.length; + } + } + } else if (element.table) { + for (const row of element.table.tableRows || []) { + for (const cell of row.tableCells || []) { + matches.push(...this._findTextMatches(cell.content, findText)); + } + } + } + } + return matches; + } + + private _collectParagraphs( + content: docs_v1.Schema$StructuralElement[] | undefined, + ): StructuredParagraph[] { + const paragraphs: StructuredParagraph[] = []; + for (const element of content || []) { + if ( + element.paragraph && + element.startIndex !== undefined && + element.endIndex !== undefined + ) { + paragraphs.push({ + startIndex: element.startIndex, + endIndex: element.endIndex, + style: + element.paragraph.paragraphStyle?.namedStyleType || 'NORMAL_TEXT', + text: this._readStructuralElement(element), + }); + } else if (element.table) { + for (const row of element.table.tableRows || []) { + for (const cell of row.tableCells || []) { + paragraphs.push(...this._collectParagraphs(cell.content)); + } + } + } + } + return paragraphs; + } + + private async _loadStructuredDocument(documentId: string) { + const id = validateAndExtractDocId(documentId); + const docs = await this.getDocsClient(); + const response = await docs.documents.get({ + documentId: id, + fields: SAFE_DOCS_FIELD_MASK, + includeTabsContent: true, + suggestionsViewMode: 'PREVIEW_WITHOUT_SUGGESTIONS', + }); + const revisionId = response.data.revisionId; + if (!revisionId) { + throw new Error('Google Docs response did not include a revision ID.'); + } + return { + id, + docs, + revisionId, + tabs: this._flattenTabs(response.data.tabs || []), + }; + } + + private _resolveSection( + tabs: docs_v1.Schema$Tab[], + tabId: string, + startHeading: string, + endHeading?: string, + ): SectionLocation { + const tab = tabs.find( + (candidate) => candidate.tabProperties?.tabId === tabId, + ); + if (!tab) { + throw new Error(`Tab with ID ${tabId} not found.`); + } + const content = tab.documentTab?.body?.content || []; + const headings = content.filter((element) => { + const style = element.paragraph?.paragraphStyle?.namedStyleType || ''; + return element.paragraph && style.startsWith('HEADING_'); + }); + const startMatches = headings.filter( + (element) => this._readStructuralElement(element).trim() === startHeading, + ); + if (startMatches.length !== 1) { + throw new Error( + `Expected exactly one heading "${startHeading}" in tab ${tabId}; found ${startMatches.length}.`, + ); + } + const startIndex = startMatches[0].startIndex; + if (startIndex === undefined) { + throw new Error(`Heading "${startHeading}" has no start index.`); + } + + let endIndex: number; + if (endHeading) { + const endMatches = headings.filter( + (element) => + (element.startIndex ?? 0) > startIndex && + this._readStructuralElement(element).trim() === endHeading, + ); + if (endMatches.length !== 1 || endMatches[0].startIndex === undefined) { + throw new Error( + `Expected exactly one heading "${endHeading}" after "${startHeading}" in tab ${tabId}; found ${endMatches.length}.`, + ); + } + endIndex = endMatches[0].startIndex; + } else { + const finalEndIndex = content[content.length - 1]?.endIndex; + if (finalEndIndex === undefined) { + throw new Error(`Tab ${tabId} has no editable body content.`); + } + endIndex = Math.max(startIndex, finalEndIndex - 1); + } + const currentText = content + .filter( + (element) => + (element.startIndex ?? -1) >= startIndex && + (element.startIndex ?? endIndex) < endIndex, + ) + .map((element) => this._readStructuralElement(element)) + .join(''); + return { + range: { tabId, startIndex, endIndex }, + currentText, + }; + } + + private _renderBlocks( + blocks: DocsReplacementBlock[], + startIndex: number, + tabId: string, + ) { let text = ''; - if (content) { - content.forEach((element) => { - text += this._readStructuralElement(element); - }); + const styleRequests: docs_v1.Schema$UpdateParagraphStyleRequest[] = []; + for (const block of blocks) { + const blockText = block.text.endsWith('\n') + ? block.text + : `${block.text}\n`; + const blockStart = startIndex + text.length; + text += blockText; + if (block.style) { + const namedStyleType = + DocsService.HEADING_STYLES[block.style.toLowerCase()] || block.style; + if ( + !/^HEADING_[1-6]$/.test(namedStyleType) && + namedStyleType !== 'NORMAL_TEXT' + ) { + throw new Error(`Unsupported paragraph style: ${block.style}`); + } + styleRequests.push({ + range: { + tabId, + startIndex: blockStart, + endIndex: blockStart + blockText.length, + }, + paragraphStyle: { namedStyleType }, + fields: 'namedStyleType', + }); + } } - return text; + return { text, styleRequests }; + } + + private _docsError(tool: string, error: unknown) { + const message = error instanceof Error ? error.message : String(error); + logToFile(`[DocsService] Error during ${tool}: ${message}`); + return { + isError: true, + content: [ + { type: 'text' as const, text: JSON.stringify({ error: message }) }, + ], + }; } }