1  import type { TextBlock } from '@anthropic-ai/sdk/resources/index.mjs'
  2  import { Box } from 'ink'
  3  import * as React from 'react'
  4  import { z } from 'zod'
  5  import type { Tool } from '../../Tool.js'
  6  import { FallbackToolUseRejectedMessage } from '../../components/FallbackToolUseRejectedMessage.js'
  7  import { HighlightedCode } from '../../components/HighlightedCode.js'
  8  import { getContext } from '../../context.js'
  9  import { Message, query } from '../../query.js'
 10  import { lastX } from '../../utils/generators.js'
 11  import { createUserMessage } from '../../utils/messages.js'
 12  import { BashTool } from '../BashTool/BashTool.js'
 13  import { FileReadTool } from '../FileReadTool/FileReadTool.js'
 14  import { FileWriteTool } from '../FileWriteTool/FileWriteTool.js'
 15  import { GlobTool } from '../GlobTool/GlobTool.js'
 16  import { GrepTool } from '../GrepTool/GrepTool.js'
 17  import { LSTool } from '../lsTool/lsTool.js'
 18  import { ARCHITECT_SYSTEM_PROMPT, DESCRIPTION } from './prompt.js'
 19  
 20  const FS_EXPLORATION_TOOLS: Tool[] = [
 21    BashTool,
 22    LSTool,
 23    FileReadTool,
 24    FileWriteTool,
 25    GlobTool,
 26    GrepTool,
 27  ]
 28  
 29  const inputSchema = z.strictObject({
 30    prompt: z
 31      .string()
 32      .describe('The technical request or coding task to analyze'),
 33    context: z
 34      .string()
 35      .describe('Optional context from previous conversation or system state')
 36      .optional(),
 37  })
 38  
 39  export const ArchitectTool = {
 40    name: 'Architect',
 41    async description() {
 42      return DESCRIPTION
 43    },
 44    inputSchema,
 45    isReadOnly() {
 46      return true
 47    },
 48    userFacingName() {
 49      return 'Architect'
 50    },
 51    async isEnabled() {
 52      return false
 53    },
 54    needsPermissions() {
 55      return false
 56    },
 57    async *call({ prompt, context }, toolUseContext, canUseTool) {
 58      const content = context
 59        ? `<context>${context}</context>\n\n${prompt}`
 60        : prompt
 61  
 62      const userMessage = createUserMessage(content)
 63  
 64      const messages: Message[] = [userMessage]
 65  
 66      // We only allow the file exploration tools to be used in the architect tool
 67      const allowedTools = (toolUseContext.options.tools ?? []).filter(_ =>
 68        FS_EXPLORATION_TOOLS.map(_ => _.name).includes(_.name),
 69      )
 70  
 71      const lastResponse = await lastX(
 72        query(
 73          messages,
 74          [ARCHITECT_SYSTEM_PROMPT],
 75          await getContext(),
 76          canUseTool,
 77          {
 78            ...toolUseContext,
 79            options: { ...toolUseContext.options, tools: allowedTools },
 80          },
 81        ),
 82      )
 83  
 84      if (lastResponse.type !== 'assistant') {
 85        throw new Error('Invalid response from Claude API')
 86      }
 87  
 88      const data = lastResponse.message.content.filter(_ => _.type === 'text')
 89      yield {
 90        type: 'result',
 91        data,
 92        resultForAssistant: this.renderResultForAssistant(data),
 93      }
 94    },
 95    async prompt() {
 96      return DESCRIPTION
 97    },
 98    renderResultForAssistant(data) {
 99      return data
100    },
101    renderToolUseMessage(input) {
102      return Object.entries(input)
103        .map(([key, value]) => `${key}: ${JSON.stringify(value)}`)
104        .join(', ')
105    },
106    renderToolResultMessage(content) {
107      return (
108        <Box flexDirection="column" gap={1}>
109          <HighlightedCode
110            code={content.map(_ => _.text).join('\n')}
111            language="markdown"
112          />
113        </Box>
114      )
115    },
116    renderToolUseRejectedMessage() {
117      return <FallbackToolUseRejectedMessage />
118    },
119  } satisfies Tool<typeof inputSchema, TextBlock[]>