Merge pull request #599 from buster-so/dallin/bus-1485-clarifications-and-todolist-runtime-context

Dallin/bus-1485-clarifications-and-todolist-runtime-context

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<typeof ClarifyingQuestionSchema>;
+
+export type DocsAgentContext = {
+  sandbox: Sandbox;
+  todoList: string;
+  clarificationQuestion: MessageUserClarifyingQuestion;
+};

packages/ai/src/context/sandbox-context.ts

@@ -1,9 +0,0 @@
-import type { Sandbox } from '@buster/sandbox';
-
-export enum SandboxContextKey {
-  Sandbox = 'sandbox',
-}
-
-export type SandboxContext = {
-  sandbox: Sandbox;
-};

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

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 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<SandboxContext>;
+  let runtimeContext: RuntimeContext<DocsAgentContext>;
 
   beforeEach(() => {
     vi.clearAllMocks();
-    runtimeContext = new RuntimeContext<SandboxContext>();
+    runtimeContext = new RuntimeContext<DocsAgentContext>();
   });
 
   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"' }],

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 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<typeof inputSchema>,
-    runtimeContext: RuntimeContext<SandboxContext>
+    runtimeContext: RuntimeContext<DocsAgentContext>
   ): Promise<z.infer<typeof outputSchema>> => {
     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<typeof inputSchema>;
-    runtimeContext: RuntimeContext<SandboxContext>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
   }) => {
     return await executeBashCommands(context, runtimeContext);
   },

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<SandboxContext>;
+  let runtimeContext: RuntimeContext<DocsAgentContext>;
 
   beforeEach(() => {
     vi.clearAllMocks();
-    runtimeContext = new RuntimeContext<SandboxContext>();
+    runtimeContext = new RuntimeContext<DocsAgentContext>();
   });
 
   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' }],

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<typeof createFilesInputSchema>,
-    runtimeContext: RuntimeContext<SandboxContext>
+    runtimeContext: RuntimeContext<DocsAgentContext>
   ): Promise<z.infer<typeof createFilesOutputSchema>> => {
     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<typeof createFilesInputSchema>;
-    runtimeContext: RuntimeContext<SandboxContext>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
   }) => {
     return await createFilesExecution(context, runtimeContext);
   },

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<SandboxContext>;
+  let runtimeContext: RuntimeContext<DocsAgentContext>;
 
   beforeEach(() => {
     vi.clearAllMocks();
-    runtimeContext = new RuntimeContext<SandboxContext>();
+    runtimeContext = new RuntimeContext<DocsAgentContext>();
   });
 
   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' }],

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<typeof deleteFilesInputSchema>,
-    runtimeContext: RuntimeContext<SandboxContext>
+    runtimeContext: RuntimeContext<DocsAgentContext>
   ): Promise<z.infer<typeof deleteFilesOutputSchema>> => {
     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<typeof deleteFilesInputSchema>;
-    runtimeContext: RuntimeContext<SandboxContext>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
   }) => {
     return await deleteFilesExecution(context, runtimeContext);
   },

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<typeof editFilesInputSchema>,
-    runtimeContext: RuntimeContext<SandboxContext>
+    runtimeContext: RuntimeContext<DocsAgentContext>
   ): Promise<z.infer<typeof editFilesOutputSchema>> => {
     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<typeof editFilesInputSchema>;
-    runtimeContext: RuntimeContext<SandboxContext>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
   }) => {
     return await editFilesExecution(context, runtimeContext);
   },

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: {

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<typeof grepSearchOutputSchema>;
 const grepSearchExecution = wrapTraced(
   async (
     params: z.infer<typeof grepSearchInputSchema>,
-    runtimeContext: RuntimeContext<SandboxContext>
+    runtimeContext: RuntimeContext<DocsAgentContext>
   ): Promise<z.infer<typeof grepSearchOutputSchema>> => {
     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<typeof grepSearchInputSchema>;
-    runtimeContext: RuntimeContext<SandboxContext>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
   }) => {
     return await grepSearchExecution(context, runtimeContext);
   },

packages/ai/src/tools/file-tools/index.ts

@@ -1 +1 @@
-export { bashExecute } from './bash-execute-tool';
+export { bashExecute } from './bash-tool/bash-execute-tool';

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({

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<typeof lsFilesInputSchema>,
-    runtimeContext: RuntimeContext<SandboxContext>
+    runtimeContext: RuntimeContext<DocsAgentContext>
   ): Promise<z.infer<typeof lsFilesOutputSchema>> => {
     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<typeof lsFilesInputSchema>;
-    runtimeContext: RuntimeContext<SandboxContext>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
   }) => {
     return await lsFilesExecution(context, runtimeContext);
   },

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: {

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<typeof readFilesInputSchema>,
-    runtimeContext: RuntimeContext<SandboxContext>
+    runtimeContext: RuntimeContext<DocsAgentContext>
   ): Promise<z.infer<typeof readFilesOutputSchema>> => {
     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<typeof readFilesInputSchema>;
-    runtimeContext: RuntimeContext<SandboxContext>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
   }) => {
     return await readFilesExecution(context, runtimeContext);
   },

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';

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<DocsAgentContext>;
+
+  beforeEach(() => {
+    runtimeContext = new RuntimeContext<DocsAgentContext>();
+  });
+
+  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);
+  });
+});

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<typeof checkOffTodoListInputSchema>,
+    runtimeContext: RuntimeContext<DocsAgentContext>
+  ): Promise<z.infer<typeof checkOffTodoListOutputSchema>> => {
+    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<typeof checkOffTodoListInputSchema>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
+  }) => {
+    return await checkOffTodoListExecution(context, runtimeContext);
+  },
+});
+
+export default checkOffTodoList;

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<DocsAgentContext>;
+
+  beforeEach(() => {
+    runtimeContext = new RuntimeContext<DocsAgentContext>();
+  });
+
+  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: <>?/@#$%',
+    });
+  });
+});

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<typeof updateClarificationsInputSchema>,
+    runtimeContext: RuntimeContext<DocsAgentContext>
+  ): Promise<z.infer<typeof updateClarificationsOutputSchema>> => {
+    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<typeof updateClarificationsInputSchema>;
+    runtimeContext: RuntimeContext<DocsAgentContext>;
+  }) => {
+    return await updateClarificationsExecution(context, runtimeContext);
+  },
+});
+
+export default updateClarificationsFile;