From 65a39041a813ef25da7a2132e8e0942c659e7cd4 Mon Sep 17 00:00:00 2001 From: dal Date: Tue, 22 Jul 2025 12:14:43 -0600 Subject: [PATCH 1/3] renamed the sandbox context to the docs agent context and added in the clarification and the todos list --- packages/ai/src/context/docs-agent-context.ts | 22 ++++++++++++ packages/ai/src/context/sandbox-context.ts | 9 ----- packages/ai/src/tools/file-tools/CLAUDE.md | 4 +-- .../file-tools/bash-execute-tool.test.ts | 12 +++---- .../src/tools/file-tools/bash-execute-tool.ts | 8 ++--- .../create-file-tool.test.ts | 12 +++---- .../create-files-tool/create-file-tool.ts | 8 ++--- .../delete-files-tool.test.ts | 12 +++---- .../delete-files-tool/delete-files-tool.ts | 8 ++--- .../edit-files-tool/edit-files-tool.ts | 8 ++--- .../grep-search-tool.int.test.ts | 8 ++--- .../grep-search-tool/grep-search-tool.ts | 8 ++--- .../ls-files-tool/ls-files-tool.test.ts | 6 ++-- .../file-tools/ls-files-tool/ls-files-tool.ts | 8 ++--- .../read-files-tool.int.test.ts | 8 ++--- .../read-files-tool/read-files-tool.ts | 8 ++--- .../modify-dashboards-file-tool.ts | 35 ++++++++++++------- .../utils/retry/retry-mechanism.int.test.ts | 4 ++- 18 files changed, 106 insertions(+), 82 deletions(-) create mode 100644 packages/ai/src/context/docs-agent-context.ts delete mode 100644 packages/ai/src/context/sandbox-context.ts diff --git a/packages/ai/src/context/docs-agent-context.ts b/packages/ai/src/context/docs-agent-context.ts new file mode 100644 index 000000000..1a9311992 --- /dev/null +++ b/packages/ai/src/context/docs-agent-context.ts @@ -0,0 +1,22 @@ +import type { Sandbox } from '@buster/sandbox'; +import { z } from 'zod'; + +export enum DocsAgentContextKey { + Sandbox = 'sandbox', + TodoListFile = 'todoListFile', + ClarificationFile = 'clarificationFile', +} + +export const ClarifyingQuestionSchema = z.object({ + issue: z.string(), + context: z.string(), + clarificationQuestion: z.string(), +}); + +export type MessageUserClarifyingQuestion = z.infer; + +export type DocsAgentContext = { + sandbox: Sandbox; + todoList: string; + clarificationQuestion: MessageUserClarifyingQuestion; +}; diff --git a/packages/ai/src/context/sandbox-context.ts b/packages/ai/src/context/sandbox-context.ts deleted file mode 100644 index 12398e2fa..000000000 --- a/packages/ai/src/context/sandbox-context.ts +++ /dev/null @@ -1,9 +0,0 @@ -import type { Sandbox } from '@buster/sandbox'; - -export enum SandboxContextKey { - Sandbox = 'sandbox', -} - -export type SandboxContext = { - sandbox: Sandbox; -}; diff --git a/packages/ai/src/tools/file-tools/CLAUDE.md b/packages/ai/src/tools/file-tools/CLAUDE.md index 21093701b..87b6300a1 100644 --- a/packages/ai/src/tools/file-tools/CLAUDE.md +++ b/packages/ai/src/tools/file-tools/CLAUDE.md @@ -86,11 +86,11 @@ const results = await readFiles(['file.txt']); Tools should check for sandbox in runtime context and adapt accordingly: ```typescript -import { SandboxContextKey } from '@buster/ai/context/sandbox-context'; +import { DocsAgentContextKey } from '@buster/ai/context/docs-agent-context'; import { runTypescript } from '@buster/sandbox'; // In your tool execution: -const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); +const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { // Generate CommonJS/sync code for sandbox diff --git a/packages/ai/src/tools/file-tools/bash-execute-tool.test.ts b/packages/ai/src/tools/file-tools/bash-execute-tool.test.ts index 472f4fbd5..b9caebdea 100644 --- a/packages/ai/src/tools/file-tools/bash-execute-tool.test.ts +++ b/packages/ai/src/tools/file-tools/bash-execute-tool.test.ts @@ -1,7 +1,7 @@ import { RuntimeContext } from '@mastra/core/runtime-context'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../context/docs-agent-context'; import { bashExecute } from './bash-execute-tool'; vi.mock('@buster/sandbox', () => ({ @@ -21,11 +21,11 @@ const mockGenerateBashExecuteCode = vi.mocked(generateBashExecuteCode); const mockExecuteBashCommandsSafely = vi.mocked(executeBashCommandsSafely); describe('bash-execute-tool', () => { - let runtimeContext: RuntimeContext; + let runtimeContext: RuntimeContext; beforeEach(() => { vi.clearAllMocks(); - runtimeContext = new RuntimeContext(); + runtimeContext = new RuntimeContext(); }); afterEach(() => { @@ -53,7 +53,7 @@ describe('bash-execute-tool', () => { it('should execute with sandbox when available', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { commands: [{ command: 'echo "hello"' }], @@ -125,7 +125,7 @@ describe('bash-execute-tool', () => { it('should handle sandbox execution errors', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { commands: [{ command: 'echo "hello"' }], @@ -187,7 +187,7 @@ describe('bash-execute-tool', () => { it('should handle JSON parse errors from sandbox', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { commands: [{ command: 'echo "hello"' }], diff --git a/packages/ai/src/tools/file-tools/bash-execute-tool.ts b/packages/ai/src/tools/file-tools/bash-execute-tool.ts index c4cd21dc7..ca6461f45 100644 --- a/packages/ai/src/tools/file-tools/bash-execute-tool.ts +++ b/packages/ai/src/tools/file-tools/bash-execute-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../context/docs-agent-context'; const bashCommandSchema = z.object({ command: z.string().describe('The bash command to execute'), @@ -33,7 +33,7 @@ const outputSchema = z.object({ const executeBashCommands = wrapTraced( async ( input: z.infer, - runtimeContext: RuntimeContext + runtimeContext: RuntimeContext ): Promise> => { const commands = Array.isArray(input.commands) ? input.commands : [input.commands]; @@ -43,7 +43,7 @@ const executeBashCommands = wrapTraced( try { // Check if sandbox is available in runtime context - const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); + const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { const { generateBashExecuteCode } = await import('./bash-execute-functions'); @@ -106,7 +106,7 @@ export const bashExecute = createTool({ runtimeContext, }: { context: z.infer; - runtimeContext: RuntimeContext; + runtimeContext: RuntimeContext; }) => { return await executeBashCommands(context, runtimeContext); }, diff --git a/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.test.ts b/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.test.ts index 2be178dee..36995d524 100644 --- a/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.test.ts +++ b/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.test.ts @@ -1,7 +1,7 @@ import { RuntimeContext } from '@mastra/core/runtime-context'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; import { createFiles } from './create-file-tool'; vi.mock('@buster/sandbox', () => ({ @@ -21,11 +21,11 @@ const mockGenerateFileCreateCode = vi.mocked(generateFileCreateCode); const mockCreateFilesSafely = vi.mocked(createFilesSafely); describe('create-file-tool', () => { - let runtimeContext: RuntimeContext; + let runtimeContext: RuntimeContext; beforeEach(() => { vi.clearAllMocks(); - runtimeContext = new RuntimeContext(); + runtimeContext = new RuntimeContext(); }); afterEach(() => { @@ -63,7 +63,7 @@ describe('create-file-tool', () => { it('should execute with sandbox when available', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { files: [{ path: '/test/file.txt', content: 'test content' }], @@ -117,7 +117,7 @@ describe('create-file-tool', () => { it('should handle sandbox execution errors', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { files: [{ path: '/test/file.txt', content: 'test content' }], @@ -191,7 +191,7 @@ describe('create-file-tool', () => { it('should handle JSON parse errors from sandbox', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { files: [{ path: '/test/file.txt', content: 'test content' }], diff --git a/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.ts b/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.ts index 80be8c2cf..14c3a1514 100644 --- a/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.ts +++ b/packages/ai/src/tools/file-tools/create-files-tool/create-file-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; const fileCreateParamsSchema = z.object({ path: z.string().describe('The relative or absolute path to create the file at'), @@ -33,7 +33,7 @@ const createFilesOutputSchema = z.object({ const createFilesExecution = wrapTraced( async ( params: z.infer, - runtimeContext: RuntimeContext + runtimeContext: RuntimeContext ): Promise> => { const { files } = params; @@ -43,7 +43,7 @@ const createFilesExecution = wrapTraced( try { // Check if sandbox is available in runtime context - const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); + const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { // Execute in sandbox @@ -132,7 +132,7 @@ export const createFiles = createTool({ runtimeContext, }: { context: z.infer; - runtimeContext: RuntimeContext; + runtimeContext: RuntimeContext; }) => { return await createFilesExecution(context, runtimeContext); }, diff --git a/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.test.ts b/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.test.ts index afd982670..a8d1a5b14 100644 --- a/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.test.ts +++ b/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.test.ts @@ -1,7 +1,7 @@ import { RuntimeContext } from '@mastra/core/runtime-context'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; import { deleteFiles } from './delete-files-tool'; vi.mock('@buster/sandbox', () => ({ @@ -21,11 +21,11 @@ const mockGenerateFileDeleteCode = vi.mocked(generateFileDeleteCode); const mockDeleteFilesSafely = vi.mocked(deleteFilesSafely); describe('delete-files-tool', () => { - let runtimeContext: RuntimeContext; + let runtimeContext: RuntimeContext; beforeEach(() => { vi.clearAllMocks(); - runtimeContext = new RuntimeContext(); + runtimeContext = new RuntimeContext(); }); afterEach(() => { @@ -58,7 +58,7 @@ describe('delete-files-tool', () => { it('should execute with sandbox when available', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { files: [{ path: '/test/file.txt' }], @@ -106,7 +106,7 @@ describe('delete-files-tool', () => { it('should handle sandbox execution errors', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { files: [{ path: '/test/file.txt' }], @@ -175,7 +175,7 @@ describe('delete-files-tool', () => { it('should handle JSON parse errors from sandbox', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox as any); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox as any); const input = { files: [{ path: '/test/file.txt' }], diff --git a/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.ts b/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.ts index a574c0c74..4119029d5 100644 --- a/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.ts +++ b/packages/ai/src/tools/file-tools/delete-files-tool/delete-files-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; const deleteFilesInputSchema = z.object({ files: z @@ -28,7 +28,7 @@ const deleteFilesOutputSchema = z.object({ const deleteFilesExecution = wrapTraced( async ( params: z.infer, - runtimeContext: RuntimeContext + runtimeContext: RuntimeContext ): Promise> => { const { files } = params; @@ -37,7 +37,7 @@ const deleteFilesExecution = wrapTraced( } try { - const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); + const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { const { generateFileDeleteCode } = await import('./delete-files-functions'); @@ -123,7 +123,7 @@ export const deleteFiles = createTool({ runtimeContext, }: { context: z.infer; - runtimeContext: RuntimeContext; + runtimeContext: RuntimeContext; }) => { return await deleteFilesExecution(context, runtimeContext); }, diff --git a/packages/ai/src/tools/file-tools/edit-files-tool/edit-files-tool.ts b/packages/ai/src/tools/file-tools/edit-files-tool/edit-files-tool.ts index 25f0b8635..1c5f3d828 100644 --- a/packages/ai/src/tools/file-tools/edit-files-tool/edit-files-tool.ts +++ b/packages/ai/src/tools/file-tools/edit-files-tool/edit-files-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; const editFileParamsSchema = z.object({ filePath: z.string().describe('Relative or absolute path to the file'), @@ -46,7 +46,7 @@ const editFilesOutputSchema = z.object({ const editFilesExecution = wrapTraced( async ( params: z.infer, - runtimeContext: RuntimeContext + runtimeContext: RuntimeContext ): Promise> => { const { edits } = params; @@ -58,7 +58,7 @@ const editFilesExecution = wrapTraced( } try { - const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); + const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { const { generateFileEditCode } = await import('./edit-files'); @@ -182,7 +182,7 @@ For bulk operations, each edit is processed independently and the tool returns b runtimeContext, }: { context: z.infer; - runtimeContext: RuntimeContext; + runtimeContext: RuntimeContext; }) => { return await editFilesExecution(context, runtimeContext); }, diff --git a/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.int.test.ts b/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.int.test.ts index ca54aa1e1..b9a70f8f5 100644 --- a/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.int.test.ts +++ b/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.int.test.ts @@ -1,7 +1,7 @@ import { type Sandbox, createSandbox } from '@buster/sandbox'; import { RuntimeContext } from '@mastra/core/runtime-context'; import { afterAll, beforeAll, describe, expect, it } from 'vitest'; -import { SandboxContextKey } from '../../../context/sandbox-context'; +import { DocsAgentContextKey } from '../../../context/docs-agent-context'; import { grepSearch } from './grep-search-tool'; describe('grep-search-tool integration test', () => { @@ -36,7 +36,7 @@ describe('grep-search-tool integration test', () => { await sandbox.process.codeRun(createFilesCode); const runtimeContext = new RuntimeContext(); - runtimeContext.set(SandboxContextKey.Sandbox, sandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, sandbox); const result = await grepSearch.execute({ context: { @@ -95,7 +95,7 @@ describe('grep-search-tool integration test', () => { it.skipIf(!hasApiKey)('should handle non-existent files in sandbox', async () => { const runtimeContext = new RuntimeContext(); - runtimeContext.set(SandboxContextKey.Sandbox, sandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, sandbox); const result = await grepSearch.execute({ context: { @@ -132,7 +132,7 @@ describe('grep-search-tool integration test', () => { await sandbox.process.codeRun(createFileCode); const runtimeContext = new RuntimeContext(); - runtimeContext.set(SandboxContextKey.Sandbox, sandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, sandbox); const result = await grepSearch.execute({ context: { diff --git a/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.ts b/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.ts index b92b1cba4..6b01062a4 100644 --- a/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.ts +++ b/packages/ai/src/tools/file-tools/grep-search-tool/grep-search-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; const grepSearchConfigSchema = z .object({ @@ -87,7 +87,7 @@ export type GrepSearchOutput = z.infer; const grepSearchExecution = wrapTraced( async ( params: z.infer, - runtimeContext: RuntimeContext + runtimeContext: RuntimeContext ): Promise> => { const { searches: rawSearches } = params; @@ -104,7 +104,7 @@ const grepSearchExecution = wrapTraced( } try { - const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); + const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { const { generateGrepSearchCode } = await import('./grep-search'); @@ -199,7 +199,7 @@ export const grepSearch = createTool({ runtimeContext, }: { context: z.infer; - runtimeContext: RuntimeContext; + runtimeContext: RuntimeContext; }) => { return await grepSearchExecution(context, runtimeContext); }, diff --git a/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.test.ts b/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.test.ts index 515f6f544..d019c506b 100644 --- a/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.test.ts +++ b/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.test.ts @@ -1,6 +1,6 @@ import { RuntimeContext } from '@mastra/core/runtime-context'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { SandboxContextKey } from '../../../context/sandbox-context'; +import { DocsAgentContextKey } from '../../../context/docs-agent-context'; const mockRunTypescript = vi.fn(); const mockLsFilesSafely = vi.fn(); @@ -70,7 +70,7 @@ describe('ls-files-tool', () => { it('should execute with sandbox when available', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox); mockGenerateLsCode.mockReturnValue('generated code'); mockRunTypescript.mockResolvedValue({ @@ -116,7 +116,7 @@ describe('ls-files-tool', () => { it('should handle sandbox execution failure', async () => { const mockSandbox = { process: { codeRun: vi.fn() } }; - runtimeContext.set(SandboxContextKey.Sandbox, mockSandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, mockSandbox); mockGenerateLsCode.mockReturnValue('generated code'); mockRunTypescript.mockResolvedValue({ diff --git a/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.ts b/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.ts index f9f5dd62b..28d5f47cf 100644 --- a/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.ts +++ b/packages/ai/src/tools/file-tools/ls-files-tool/ls-files-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; import type { LsOptions } from './ls-files-impl'; const lsOptionsSchema = z.object({ @@ -63,7 +63,7 @@ const lsFilesOutputSchema = z.object({ const lsFilesExecution = wrapTraced( async ( params: z.infer, - runtimeContext: RuntimeContext + runtimeContext: RuntimeContext ): Promise> => { const { paths, options } = params; @@ -72,7 +72,7 @@ const lsFilesExecution = wrapTraced( } try { - const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); + const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { const { generateLsCode } = await import('./ls-files-impl'); @@ -185,7 +185,7 @@ export const lsFiles = createTool({ runtimeContext, }: { context: z.infer; - runtimeContext: RuntimeContext; + runtimeContext: RuntimeContext; }) => { return await lsFilesExecution(context, runtimeContext); }, diff --git a/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.int.test.ts b/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.int.test.ts index ffac13172..d5365ccb8 100644 --- a/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.int.test.ts +++ b/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.int.test.ts @@ -1,7 +1,7 @@ import { type Sandbox, createSandbox } from '@buster/sandbox'; import { RuntimeContext } from '@mastra/core/runtime-context'; import { afterAll, beforeAll, describe, expect, it } from 'vitest'; -import { SandboxContextKey } from '../../../context/sandbox-context'; +import { DocsAgentContextKey } from '../../../context/docs-agent-context'; import { readFiles } from './read-files-tool'; describe('read-files-tool integration test', () => { @@ -40,7 +40,7 @@ describe('read-files-tool integration test', () => { // Now test reading files with the tool const runtimeContext = new RuntimeContext(); - runtimeContext.set(SandboxContextKey.Sandbox, sandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, sandbox); const result = await readFiles.execute({ context: { @@ -66,7 +66,7 @@ describe('read-files-tool integration test', () => { it.skipIf(!hasApiKey)('should handle non-existent files in sandbox', async () => { const runtimeContext = new RuntimeContext(); - runtimeContext.set(SandboxContextKey.Sandbox, sandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, sandbox); const result = await readFiles.execute({ context: { @@ -97,7 +97,7 @@ describe('read-files-tool integration test', () => { await sandbox.process.codeRun(createFilesCode); const runtimeContext = new RuntimeContext(); - runtimeContext.set(SandboxContextKey.Sandbox, sandbox); + runtimeContext.set(DocsAgentContextKey.Sandbox, sandbox); const result = await readFiles.execute({ context: { diff --git a/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.ts b/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.ts index 748b77481..8b8fcb240 100644 --- a/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.ts +++ b/packages/ai/src/tools/file-tools/read-files-tool/read-files-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type SandboxContext, SandboxContextKey } from '../../../context/sandbox-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; import type { AnalystRuntimeContext } from '../../../schemas/workflow-schemas'; const readFilesInputSchema = z.object({ @@ -37,7 +37,7 @@ const readFilesOutputSchema = z.object({ const readFilesExecution = wrapTraced( async ( params: z.infer, - runtimeContext: RuntimeContext + runtimeContext: RuntimeContext ): Promise> => { const { files } = params; @@ -47,7 +47,7 @@ const readFilesExecution = wrapTraced( try { // Check if sandbox is available in runtime context - const sandbox = runtimeContext.get(SandboxContextKey.Sandbox); + const sandbox = runtimeContext.get(DocsAgentContextKey.Sandbox); if (sandbox) { // Execute in sandbox @@ -141,7 +141,7 @@ export const readFiles = createTool({ runtimeContext, }: { context: z.infer; - runtimeContext: RuntimeContext; + runtimeContext: RuntimeContext; }) => { return await readFilesExecution(context, runtimeContext); }, diff --git a/packages/ai/src/tools/visualization-tools/modify-dashboards-file-tool.ts b/packages/ai/src/tools/visualization-tools/modify-dashboards-file-tool.ts index 31fb92897..580970c6a 100644 --- a/packages/ai/src/tools/visualization-tools/modify-dashboards-file-tool.ts +++ b/packages/ai/src/tools/visualization-tools/modify-dashboards-file-tool.ts @@ -371,8 +371,10 @@ const modifyDashboardFiles = wrapTraced( for (const file of dashboardFilesToUpdate) { // Get current metric IDs from updated dashboard content - const newMetricIds = (file.content as DashboardYml).rows.flatMap(row => row.items).map(item => item.id); - + const newMetricIds = (file.content as DashboardYml).rows + .flatMap((row) => row.items) + .map((item) => item.id); + const existingAssociations = await tx .select({ metricFileId: metricFilesToDashboardFiles.metricFileId }) .from(metricFilesToDashboardFiles) @@ -383,10 +385,12 @@ const modifyDashboardFiles = wrapTraced( ) ) .execute(); - - const existingMetricIds = existingAssociations.map(a => a.metricFileId); - - const addedMetricIds = newMetricIds.filter((id: string) => !existingMetricIds.includes(id)); + + const existingMetricIds = existingAssociations.map((a) => a.metricFileId); + + const addedMetricIds = newMetricIds.filter( + (id: string) => !existingMetricIds.includes(id) + ); for (const metricId of addedMetricIds) { await tx .insert(metricFilesToDashboardFiles) @@ -396,25 +400,30 @@ const modifyDashboardFiles = wrapTraced( createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(), deletedAt: null, - createdBy: userId + createdBy: userId, }) .onConflictDoUpdate({ - target: [metricFilesToDashboardFiles.metricFileId, metricFilesToDashboardFiles.dashboardFileId], + target: [ + metricFilesToDashboardFiles.metricFileId, + metricFilesToDashboardFiles.dashboardFileId, + ], set: { deletedAt: null, - updatedAt: new Date().toISOString() - } + updatedAt: new Date().toISOString(), + }, }) .execute(); } - - const removedMetricIds = existingMetricIds.filter((id: string) => !newMetricIds.includes(id)); + + const removedMetricIds = existingMetricIds.filter( + (id: string) => !newMetricIds.includes(id) + ); if (removedMetricIds.length > 0) { await tx .update(metricFilesToDashboardFiles) .set({ deletedAt: new Date().toISOString(), - updatedAt: new Date().toISOString() + updatedAt: new Date().toISOString(), }) .where( and( diff --git a/packages/ai/src/utils/retry/retry-mechanism.int.test.ts b/packages/ai/src/utils/retry/retry-mechanism.int.test.ts index fab4c3288..c49b4d06b 100644 --- a/packages/ai/src/utils/retry/retry-mechanism.int.test.ts +++ b/packages/ai/src/utils/retry/retry-mechanism.int.test.ts @@ -220,7 +220,9 @@ describe('Retry Mechanism Integration Tests', () => { // The conversation should start with the original messages const userMessages = result.conversationHistory.filter((msg) => msg.role === 'user'); - const assistantMessages = result.conversationHistory.filter((msg) => msg.role === 'assistant'); + const assistantMessages = result.conversationHistory.filter( + (msg) => msg.role === 'assistant' + ); expect(userMessages.length).toBeGreaterThan(0); expect(assistantMessages.length).toBeGreaterThan(0); From 7337e11b2503706b020e339946ae23c99440ca48 Mon Sep 17 00:00:00 2001 From: dal Date: Tue, 22 Jul 2025 12:15:59 -0600 Subject: [PATCH 2/3] move bash into its own folder --- .../tools/file-tools/{ => bash-tool}/bash-execute-functions.ts | 0 .../tools/file-tools/{ => bash-tool}/bash-execute-tool.test.ts | 2 +- .../src/tools/file-tools/{ => bash-tool}/bash-execute-tool.ts | 2 +- packages/ai/src/tools/file-tools/index.ts | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) rename packages/ai/src/tools/file-tools/{ => bash-tool}/bash-execute-functions.ts (100%) rename packages/ai/src/tools/file-tools/{ => bash-tool}/bash-execute-tool.test.ts (99%) rename packages/ai/src/tools/file-tools/{ => bash-tool}/bash-execute-tool.ts (99%) diff --git a/packages/ai/src/tools/file-tools/bash-execute-functions.ts b/packages/ai/src/tools/file-tools/bash-tool/bash-execute-functions.ts similarity index 100% rename from packages/ai/src/tools/file-tools/bash-execute-functions.ts rename to packages/ai/src/tools/file-tools/bash-tool/bash-execute-functions.ts diff --git a/packages/ai/src/tools/file-tools/bash-execute-tool.test.ts b/packages/ai/src/tools/file-tools/bash-tool/bash-execute-tool.test.ts similarity index 99% rename from packages/ai/src/tools/file-tools/bash-execute-tool.test.ts rename to packages/ai/src/tools/file-tools/bash-tool/bash-execute-tool.test.ts index b9caebdea..3abea9d24 100644 --- a/packages/ai/src/tools/file-tools/bash-execute-tool.test.ts +++ b/packages/ai/src/tools/file-tools/bash-tool/bash-execute-tool.test.ts @@ -1,7 +1,7 @@ import { RuntimeContext } from '@mastra/core/runtime-context'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { z } from 'zod'; -import { type DocsAgentContext, DocsAgentContextKey } from '../../context/docs-agent-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; import { bashExecute } from './bash-execute-tool'; vi.mock('@buster/sandbox', () => ({ diff --git a/packages/ai/src/tools/file-tools/bash-execute-tool.ts b/packages/ai/src/tools/file-tools/bash-tool/bash-execute-tool.ts similarity index 99% rename from packages/ai/src/tools/file-tools/bash-execute-tool.ts rename to packages/ai/src/tools/file-tools/bash-tool/bash-execute-tool.ts index ca6461f45..16c6fc88c 100644 --- a/packages/ai/src/tools/file-tools/bash-execute-tool.ts +++ b/packages/ai/src/tools/file-tools/bash-tool/bash-execute-tool.ts @@ -3,7 +3,7 @@ import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; -import { type DocsAgentContext, DocsAgentContextKey } from '../../context/docs-agent-context'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; const bashCommandSchema = z.object({ command: z.string().describe('The bash command to execute'), diff --git a/packages/ai/src/tools/file-tools/index.ts b/packages/ai/src/tools/file-tools/index.ts index 8525f1a62..fb853a1ed 100644 --- a/packages/ai/src/tools/file-tools/index.ts +++ b/packages/ai/src/tools/file-tools/index.ts @@ -1 +1 @@ -export { bashExecute } from './bash-execute-tool'; +export { bashExecute } from './bash-tool/bash-execute-tool'; From a621f74910c7cfbcb08c9e808dd1352c8504e8f0 Mon Sep 17 00:00:00 2001 From: dal Date: Tue, 22 Jul 2025 13:01:01 -0600 Subject: [PATCH 3/3] feat: add new tools for managing todo lists and clarifications - Export checkOffTodoList and updateClarificationsFile from their respective modules in the planning-thinking-tools directory. --- packages/ai/src/tools/index.ts | 2 + .../check-off-todo-list-tool.test.ts | 134 +++++++++++++ .../check-off-todo-list-tool.ts | 84 ++++++++ .../update-clarifications-file-tool.test.ts | 187 ++++++++++++++++++ .../update-clarifications-file-tool.ts | 88 +++++++++ 5 files changed, 495 insertions(+) create mode 100644 packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.test.ts create mode 100644 packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.ts create mode 100644 packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.test.ts create mode 100644 packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.ts diff --git a/packages/ai/src/tools/index.ts b/packages/ai/src/tools/index.ts index cd910e734..148a77f16 100644 --- a/packages/ai/src/tools/index.ts +++ b/packages/ai/src/tools/index.ts @@ -17,3 +17,5 @@ export { lsFiles } from './file-tools/ls-files-tool/ls-files-tool'; export { grepSearch } from './file-tools/grep-search-tool/grep-search-tool'; export { bashExecute } from './file-tools'; export { deleteFiles } from './file-tools/delete-files-tool/delete-files-tool'; +export { checkOffTodoList } from './planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool'; +export { updateClarificationsFile } from './planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool'; diff --git a/packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.test.ts b/packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.test.ts new file mode 100644 index 000000000..8040cb05d --- /dev/null +++ b/packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.test.ts @@ -0,0 +1,134 @@ +import { RuntimeContext } from '@mastra/core/runtime-context'; +import { beforeEach, describe, expect, it } from 'vitest'; +import type { DocsAgentContext } from '../../../context/docs-agent-context'; +import { checkOffTodoList } from './check-off-todo-list-tool'; + +describe('checkOffTodoList', () => { + let runtimeContext: RuntimeContext; + + beforeEach(() => { + runtimeContext = new RuntimeContext(); + }); + + it('should check off a todo item successfully', async () => { + const initialTodoList = `## Todo List +- [ ] Write unit tests +- [ ] Implement feature +- [ ] Review code`; + + runtimeContext.set('todoList', initialTodoList); + + const result = await checkOffTodoList.execute({ + context: { todoItem: 'Write unit tests' }, + runtimeContext, + }); + + expect(result.success).toBe(true); + expect(result.updatedTodoList).toContain('- [x] Write unit tests'); + expect(result.updatedTodoList).toContain('- [ ] Implement feature'); + expect(result.updatedTodoList).toContain('- [ ] Review code'); + expect(result.message).toBe('Successfully checked off: "Write unit tests"'); + + // Verify context was updated + const updatedContext = runtimeContext.get('todoList'); + expect(updatedContext).toBe(result.updatedTodoList); + }); + + it('should return error when todo list is not found in context', async () => { + const result = await checkOffTodoList.execute({ + context: { todoItem: 'Some task' }, + runtimeContext, + }); + + expect(result.success).toBe(false); + expect(result.updatedTodoList).toBe(''); + expect(result.message).toBe('No todo list found in context'); + }); + + it('should return error when todo item is not found', async () => { + const todoList = `## Todo List +- [ ] Write unit tests +- [ ] Implement feature`; + + runtimeContext.set('todoList', todoList); + + const result = await checkOffTodoList.execute({ + context: { todoItem: 'Non-existent task' }, + runtimeContext, + }); + + expect(result.success).toBe(false); + expect(result.updatedTodoList).toBe(todoList); + expect(result.message).toBe( + 'Todo item "Non-existent task" not found in the list or already checked off' + ); + }); + + it('should not check off an already checked item', async () => { + const todoList = `## Todo List +- [x] Write unit tests +- [ ] Implement feature`; + + runtimeContext.set('todoList', todoList); + + const result = await checkOffTodoList.execute({ + context: { todoItem: 'Write unit tests' }, + runtimeContext, + }); + + expect(result.success).toBe(false); + expect(result.updatedTodoList).toBe(todoList); + expect(result.message).toBe( + 'Todo item "Write unit tests" not found in the list or already checked off' + ); + }); + + it('should handle first occurrence when there are duplicates', async () => { + const todoList = `## Todo List +- [ ] Write unit tests for feature A +- [ ] Write unit tests`; + + runtimeContext.set('todoList', todoList); + + const result = await checkOffTodoList.execute({ + context: { todoItem: 'Write unit tests for feature A' }, + runtimeContext, + }); + + expect(result.success).toBe(true); + expect(result.updatedTodoList).toBe(`## Todo List +- [x] Write unit tests for feature A +- [ ] Write unit tests`); + }); + + it('should validate input schema', () => { + const validInput = { todoItem: 'Test task' }; + const parsed = checkOffTodoList.inputSchema.parse(validInput); + expect(parsed).toEqual(validInput); + + expect(() => { + checkOffTodoList.inputSchema.parse({ todoItem: 123 }); + }).toThrow(); + + expect(() => { + checkOffTodoList.inputSchema.parse({}); + }).toThrow(); + }); + + it('should validate output schema', () => { + const validOutput = { + success: true, + updatedTodoList: '- [x] Done', + message: 'Success', + }; + const parsed = checkOffTodoList.outputSchema.parse(validOutput); + expect(parsed).toEqual(validOutput); + + const minimalOutput = { + success: false, + updatedTodoList: '', + }; + const minimalParsed = checkOffTodoList.outputSchema.parse(minimalOutput); + expect(minimalParsed).toEqual(minimalOutput); + }); +}); diff --git a/packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.ts b/packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.ts new file mode 100644 index 000000000..d81e931c0 --- /dev/null +++ b/packages/ai/src/tools/planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool.ts @@ -0,0 +1,84 @@ +import type { RuntimeContext } from '@mastra/core/runtime-context'; +import { createTool } from '@mastra/core/tools'; +import { wrapTraced } from 'braintrust'; +import { z } from 'zod'; +import { type DocsAgentContext, DocsAgentContextKey } from '../../../context/docs-agent-context'; + +const checkOffTodoListInputSchema = z.object({ + todoItem: z.string().describe('The exact text of the todo item to check off in the list'), +}); + +const checkOffTodoListOutputSchema = z.object({ + success: z.boolean(), + updatedTodoList: z.string().describe('The updated todo list with the item checked off'), + message: z.string().optional(), +}); + +const checkOffTodoListExecution = wrapTraced( + async ( + params: z.infer, + runtimeContext: RuntimeContext + ): Promise> => { + const { todoItem } = params; + + try { + // Get the current todo list from context + const currentTodoList = runtimeContext.get('todoList'); + + if (!currentTodoList) { + return { + success: false, + updatedTodoList: '', + message: 'No todo list found in context', + }; + } + + // Check if the item exists in the list (not already checked off) + if (!currentTodoList.includes(`- [ ] ${todoItem}`)) { + return { + success: false, + updatedTodoList: currentTodoList, + message: `Todo item "${todoItem}" not found in the list or already checked off`, + }; + } + + // Replace the unchecked item with a checked version + const updatedTodoList = currentTodoList.replace(`- [ ] ${todoItem}`, `- [x] ${todoItem}`); + + // Update the context with the new todo list + runtimeContext.set('todoList', updatedTodoList); + + return { + success: true, + updatedTodoList, + message: `Successfully checked off: "${todoItem}"`, + }; + } catch (error) { + return { + success: false, + updatedTodoList: '', + message: `Error checking off todo item: ${error instanceof Error ? error.message : 'Unknown error'}`, + }; + } + }, + { name: 'check-off-todo-list' } +); + +export const checkOffTodoList = createTool({ + id: 'check-off-todo-list', + description: + 'Check off a todo item in the todo list by replacing "- [ ]" with "- [x]". The todo list is maintained as a string in the runtime context.', + inputSchema: checkOffTodoListInputSchema, + outputSchema: checkOffTodoListOutputSchema, + execute: async ({ + context, + runtimeContext, + }: { + context: z.infer; + runtimeContext: RuntimeContext; + }) => { + return await checkOffTodoListExecution(context, runtimeContext); + }, +}); + +export default checkOffTodoList; diff --git a/packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.test.ts b/packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.test.ts new file mode 100644 index 000000000..a0f834dcd --- /dev/null +++ b/packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.test.ts @@ -0,0 +1,187 @@ +import { RuntimeContext } from '@mastra/core/runtime-context'; +import { beforeEach, describe, expect, it } from 'vitest'; +import type { DocsAgentContext } from '../../../context/docs-agent-context'; +import { updateClarificationsFile } from './update-clarifications-file-tool'; + +describe('updateClarificationsFile', () => { + let runtimeContext: RuntimeContext; + + beforeEach(() => { + runtimeContext = new RuntimeContext(); + }); + + it('should add a clarification question successfully', async () => { + const result = await updateClarificationsFile.execute({ + context: { + issue: 'Database connection configuration', + context: + 'The user mentioned they need to connect to a database but did not specify which type', + clarificationQuestion: + 'Which type of database are you using? (PostgreSQL, MySQL, MongoDB, etc.)', + }, + runtimeContext, + }); + + expect(result.success).toBe(true); + expect(result.clarification).toEqual({ + issue: 'Database connection configuration', + context: + 'The user mentioned they need to connect to a database but did not specify which type', + clarificationQuestion: + 'Which type of database are you using? (PostgreSQL, MySQL, MongoDB, etc.)', + }); + expect(result.message).toBe('Successfully added clarification question'); + + // Verify context was updated + const savedClarification = runtimeContext.get('clarificationQuestion'); + expect(savedClarification).toEqual(result.clarification); + }); + + it('should overwrite previous clarification when adding new one', async () => { + // Add first clarification + await updateClarificationsFile.execute({ + context: { + issue: 'First issue', + context: 'First context', + clarificationQuestion: 'First question?', + }, + runtimeContext, + }); + + // Add second clarification + const result = await updateClarificationsFile.execute({ + context: { + issue: 'Second issue', + context: 'Second context', + clarificationQuestion: 'Second question?', + }, + runtimeContext, + }); + + expect(result.success).toBe(true); + expect(result.clarification).toEqual({ + issue: 'Second issue', + context: 'Second context', + clarificationQuestion: 'Second question?', + }); + + // Verify only the second clarification is stored + const savedClarification = runtimeContext.get('clarificationQuestion'); + expect(savedClarification).toEqual(result.clarification); + }); + + it('should handle very long clarification content', async () => { + const longText = 'A'.repeat(1000); + + const result = await updateClarificationsFile.execute({ + context: { + issue: longText, + context: longText, + clarificationQuestion: `${longText}?`, + }, + runtimeContext, + }); + + expect(result.success).toBe(true); + expect(result.clarification?.issue).toBe(longText); + expect(result.clarification?.context).toBe(longText); + expect(result.clarification?.clarificationQuestion).toBe(`${longText}?`); + }); + + it('should validate input schema', () => { + const validInput = { + issue: 'Test issue', + context: 'Test context', + clarificationQuestion: 'Test question?', + }; + const parsed = updateClarificationsFile.inputSchema.parse(validInput); + expect(parsed).toEqual(validInput); + + // Missing required fields + expect(() => { + updateClarificationsFile.inputSchema.parse({ + issue: 'Test issue', + context: 'Test context', + }); + }).toThrow(); + + expect(() => { + updateClarificationsFile.inputSchema.parse({ + issue: 'Test issue', + clarificationQuestion: 'Test question?', + }); + }).toThrow(); + + expect(() => { + updateClarificationsFile.inputSchema.parse({ + context: 'Test context', + clarificationQuestion: 'Test question?', + }); + }).toThrow(); + + // Wrong types + expect(() => { + updateClarificationsFile.inputSchema.parse({ + issue: 123, + context: 'Test context', + clarificationQuestion: 'Test question?', + }); + }).toThrow(); + }); + + it('should validate output schema', () => { + const validOutput = { + success: true, + clarification: { + issue: 'Test issue', + context: 'Test context', + clarificationQuestion: 'Test question?', + }, + message: 'Success', + }; + const parsed = updateClarificationsFile.outputSchema.parse(validOutput); + expect(parsed).toEqual(validOutput); + + const minimalOutput = { + success: false, + }; + const minimalParsed = updateClarificationsFile.outputSchema.parse(minimalOutput); + expect(minimalParsed).toEqual(minimalOutput); + }); + + it('should handle empty strings', async () => { + const result = await updateClarificationsFile.execute({ + context: { + issue: '', + context: '', + clarificationQuestion: '', + }, + runtimeContext, + }); + + expect(result.success).toBe(true); + expect(result.clarification).toEqual({ + issue: '', + context: '', + clarificationQuestion: '', + }); + }); + + it('should handle special characters in clarification content', async () => { + const result = await updateClarificationsFile.execute({ + context: { + issue: 'Issue with "quotes" and \'apostrophes\'', + context: 'Context with\nnewlines\tand\ttabs', + clarificationQuestion: 'Question with émojis 🤔 and special chars: <>?/@#$%', + }, + runtimeContext, + }); + + expect(result.success).toBe(true); + expect(result.clarification).toEqual({ + issue: 'Issue with "quotes" and \'apostrophes\'', + context: 'Context with\nnewlines\tand\ttabs', + clarificationQuestion: 'Question with émojis 🤔 and special chars: <>?/@#$%', + }); + }); +}); diff --git a/packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.ts b/packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.ts new file mode 100644 index 000000000..71505fdb3 --- /dev/null +++ b/packages/ai/src/tools/planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool.ts @@ -0,0 +1,88 @@ +import type { RuntimeContext } from '@mastra/core/runtime-context'; +import { createTool } from '@mastra/core/tools'; +import { wrapTraced } from 'braintrust'; +import { z } from 'zod'; +import { + ClarifyingQuestionSchema, + type DocsAgentContext, + DocsAgentContextKey, + type MessageUserClarifyingQuestion, +} from '../../../context/docs-agent-context'; + +const updateClarificationsInputSchema = z.object({ + issue: z.string().describe('The issue or problem that needs clarification'), + context: z + .string() + .describe('The context around the issue to help understand what clarification is needed'), + clarificationQuestion: z + .string() + .describe('The specific question to ask the user for clarification'), +}); + +const updateClarificationsOutputSchema = z.object({ + success: z.boolean(), + clarification: ClarifyingQuestionSchema.optional(), + message: z.string().optional(), +}); + +const updateClarificationsExecution = wrapTraced( + async ( + params: z.infer, + runtimeContext: RuntimeContext + ): Promise> => { + const { issue, context, clarificationQuestion } = params; + + try { + // Create the new clarification question + const newClarification: MessageUserClarifyingQuestion = { + issue, + context, + clarificationQuestion, + }; + + // Validate the clarification against the schema + const validatedClarification = ClarifyingQuestionSchema.parse(newClarification); + + // Update the context with the new clarification + runtimeContext.set('clarificationQuestion', validatedClarification); + + return { + success: true, + clarification: validatedClarification, + message: 'Successfully added clarification question', + }; + } catch (error) { + if (error instanceof z.ZodError) { + return { + success: false, + message: `Validation error: ${error.errors.map((e) => e.message).join(', ')}`, + }; + } + + return { + success: false, + message: `Error adding clarification: ${error instanceof Error ? error.message : 'Unknown error'}`, + }; + } + }, + { name: 'update-clarifications-file' } +); + +export const updateClarificationsFile = createTool({ + id: 'update-clarifications-file', + description: + 'Add a new clarification question to the context. This tool helps agents request clarification from users when they encounter ambiguous or unclear requirements.', + inputSchema: updateClarificationsInputSchema, + outputSchema: updateClarificationsOutputSchema, + execute: async ({ + context, + runtimeContext, + }: { + context: z.infer; + runtimeContext: RuntimeContext; + }) => { + return await updateClarificationsExecution(context, runtimeContext); + }, +}); + +export default updateClarificationsFile;