@xortex/xcode 3.0.0

package/tools/McpAuthTool/McpAuthTool.ts

@@ -0,0 +1,215 @@
+import reject from 'lodash-es/reject.js'
+import { z } from 'zod/v4'
+import { performMCPOAuthFlow } from '../../services/mcp/auth.js'
+import {
+  clearMcpAuthCache,
+  reconnectMcpServerImpl,
+} from '../../services/mcp/client.js'
+import {
+  buildMcpToolName,
+  getMcpPrefix,
+} from '../../services/mcp/mcpStringUtils.js'
+import type {
+  McpHTTPServerConfig,
+  McpSSEServerConfig,
+  ScopedMcpServerConfig,
+} from '../../services/mcp/types.js'
+import type { Tool } from '../../Tool.js'
+import { errorMessage } from '../../utils/errors.js'
+import { lazySchema } from '../../utils/lazySchema.js'
+import { logMCPDebug, logMCPError } from '../../utils/log.js'
+import type { PermissionDecision } from '../../utils/permissions/PermissionResult.js'
+
+const inputSchema = lazySchema(() => z.object({}))
+type InputSchema = ReturnType<typeof inputSchema>
+
+export type McpAuthOutput = {
+  status: 'auth_url' | 'unsupported' | 'error'
+  message: string
+  authUrl?: string
+}
+
+function getConfigUrl(config: ScopedMcpServerConfig): string | undefined {
+  if ('url' in config) return config.url
+  return undefined
+}
+
+/**
+ * Creates a pseudo-tool for an MCP server that is installed but not
+ * authenticated. Surfaced in place of the server's real tools so the model
+ * knows the server exists and can start the OAuth flow on the user's behalf.
+ *
+ * When called, starts performMCPOAuthFlow with skipBrowserOpen and returns
+ * the authorization URL. The OAuth callback completes in the background;
+ * once it fires, reconnectMcpServerImpl runs and the server's real tools
+ * are swapped into appState.mcp.tools via the existing prefix-based
+ * replacement (useManageMCPConnections.updateServer wipes anything matching
+ * mcp__<server>__*, so this pseudo-tool is removed automatically).
+ */
+export function createMcpAuthTool(
+  serverName: string,
+  config: ScopedMcpServerConfig,
+): Tool<InputSchema, McpAuthOutput> {
+  const url = getConfigUrl(config)
+  const transport = config.type ?? 'stdio'
+  const location = url ? `${transport} at ${url}` : transport
+
+  const description =
+    `The \`${serverName}\` MCP server (${location}) is installed but requires authentication. ` +
+    `Call this tool to start the OAuth flow — you'll receive an authorization URL to share with the user. ` +
+    `Once the user completes authorization in their browser, the server's real tools will become available automatically.`
+
+  return {
+    name: buildMcpToolName(serverName, 'authenticate'),
+    isMcp: true,
+    mcpInfo: { serverName, toolName: 'authenticate' },
+    isEnabled: () => true,
+    isConcurrencySafe: () => false,
+    isReadOnly: () => false,
+    toAutoClassifierInput: () => serverName,
+    userFacingName: () => `${serverName} - authenticate (MCP)`,
+    maxResultSizeChars: 10_000,
+    renderToolUseMessage: () => `Authenticate ${serverName} MCP server`,
+    async description() {
+      return description
+    },
+    async prompt() {
+      return description
+    },
+    get inputSchema(): InputSchema {
+      return inputSchema()
+    },
+    async checkPermissions(input): Promise<PermissionDecision> {
+      return { behavior: 'allow', updatedInput: input }
+    },
+    async call(_input, context) {
+      // claude.ai connectors use a separate auth flow (handleClaudeAIAuth in
+      // MCPRemoteServerMenu) that we don't invoke programmatically here —
+      // just point the user at /mcp.
+      if (config.type === 'claudeai-proxy') {
+        return {
+          data: {
+            status: 'unsupported' as const,
+            message: `This is a claude.ai MCP connector. Ask the user to run /mcp and select "${serverName}" to authenticate.`,
+          },
+        }
+      }
+
+      // performMCPOAuthFlow only accepts sse/http. needs-auth state is only
+      // set on HTTP 401 (UnauthorizedError) so other transports shouldn't
+      // reach here, but be defensive.
+      if (config.type !== 'sse' && config.type !== 'http') {
+        return {
+          data: {
+            status: 'unsupported' as const,
+            message: `Server "${serverName}" uses ${transport} transport which does not support OAuth from this tool. Ask the user to run /mcp and authenticate manually.`,
+          },
+        }
+      }
+
+      const sseOrHttpConfig = config as (
+        | McpSSEServerConfig
+        | McpHTTPServerConfig
+      ) & { scope: ScopedMcpServerConfig['scope'] }
+
+      // Mirror cli/print.ts mcp_authenticate: start the flow, capture the
+      // URL via onAuthorizationUrl, return it immediately. The flow's
+      // Promise resolves later when the browser callback fires.
+      let resolveAuthUrl: ((url: string) => void) | undefined
+      const authUrlPromise = new Promise<string>(resolve => {
+        resolveAuthUrl = resolve
+      })
+
+      const controller = new AbortController()
+      const { setAppState } = context
+
+      const oauthPromise = performMCPOAuthFlow(
+        serverName,
+        sseOrHttpConfig,
+        u => resolveAuthUrl?.(u),
+        controller.signal,
+        { skipBrowserOpen: true },
+      )
+
+      // Background continuation: once OAuth completes, reconnect and swap
+      // the real tools into appState. Prefix-based replacement removes this
+      // pseudo-tool since it shares the mcp__<server>__ prefix.
+      void oauthPromise
+        .then(async () => {
+          clearMcpAuthCache()
+          const result = await reconnectMcpServerImpl(serverName, config)
+          const prefix = getMcpPrefix(serverName)
+          setAppState(prev => ({
+            ...prev,
+            mcp: {
+              ...prev.mcp,
+              clients: prev.mcp.clients.map(c =>
+                c.name === serverName ? result.client : c,
+              ),
+              tools: [
+                ...reject(prev.mcp.tools, t => t.name?.startsWith(prefix)),
+                ...result.tools,
+              ],
+              commands: [
+                ...reject(prev.mcp.commands, c => c.name?.startsWith(prefix)),
+                ...result.commands,
+              ],
+              resources: result.resources
+                ? { ...prev.mcp.resources, [serverName]: result.resources }
+                : prev.mcp.resources,
+            },
+          }))
+          logMCPDebug(
+            serverName,
+            `OAuth complete, reconnected with ${result.tools.length} tool(s)`,
+          )
+        })
+        .catch(err => {
+          logMCPError(
+            serverName,
+            `OAuth flow failed after tool-triggered start: ${errorMessage(err)}`,
+          )
+        })
+
+      try {
+        // Race: get the URL, or the flow completes without needing one
+        // (e.g. XAA with cached IdP token — silent auth).
+        const authUrl = await Promise.race([
+          authUrlPromise,
+          oauthPromise.then(() => null as string | null),
+        ])
+
+        if (authUrl) {
+          return {
+            data: {
+              status: 'auth_url' as const,
+              authUrl,
+              message: `Ask the user to open this URL in their browser to authorize the ${serverName} MCP server:\n\n${authUrl}\n\nOnce they complete the flow, the server's tools will become available automatically.`,
+            },
+          }
+        }
+
+        return {
+          data: {
+            status: 'auth_url' as const,
+            message: `Authentication completed silently for ${serverName}. The server's tools should now be available.`,
+          },
+        }
+      } catch (err) {
+        return {
+          data: {
+            status: 'error' as const,
+            message: `Failed to start OAuth flow for ${serverName}: ${errorMessage(err)}. Ask the user to run /mcp and authenticate manually.`,
+          },
+        }
+      }
+    },
+    mapToolResultToToolResultBlockParam(data, toolUseID) {
+      return {
+        tool_use_id: toolUseID,
+        type: 'tool_result',
+        content: data.message,
+      }
+    },
+  } satisfies Tool<InputSchema, McpAuthOutput>
+}

package/tools/NotebookEditTool/NotebookEditTool.ts

@@ -0,0 +1,490 @@
+import { feature } from 'bun:bundle'
+import { extname, isAbsolute, resolve } from 'path'
+import {
+  fileHistoryEnabled,
+  fileHistoryTrackEdit,
+} from 'src/utils/fileHistory.js'
+import { z } from 'zod/v4'
+import { buildTool, type ToolDef, type ToolUseContext } from '../../Tool.js'
+import type { NotebookCell, NotebookContent } from '../../types/notebook.js'
+import { getCwd } from '../../utils/cwd.js'
+import { isENOENT } from '../../utils/errors.js'
+import { getFileModificationTime, writeTextContent } from '../../utils/file.js'
+import { readFileSyncWithMetadata } from '../../utils/fileRead.js'
+import { safeParseJSON } from '../../utils/json.js'
+import { lazySchema } from '../../utils/lazySchema.js'
+import { parseCellId } from '../../utils/notebook.js'
+import { checkWritePermissionForTool } from '../../utils/permissions/filesystem.js'
+import type { PermissionDecision } from '../../utils/permissions/PermissionResult.js'
+import { jsonParse, jsonStringify } from '../../utils/slowOperations.js'
+import { NOTEBOOK_EDIT_TOOL_NAME } from './constants.js'
+import { DESCRIPTION, PROMPT } from './prompt.js'
+import {
+  getToolUseSummary,
+  renderToolResultMessage,
+  renderToolUseErrorMessage,
+  renderToolUseMessage,
+  renderToolUseRejectedMessage,
+} from './UI.js'
+
+export const inputSchema = lazySchema(() =>
+  z.strictObject({
+    notebook_path: z
+      .string()
+      .describe(
+        'The absolute path to the Jupyter notebook file to edit (must be absolute, not relative)',
+      ),
+    cell_id: z
+      .string()
+      .optional()
+      .describe(
+        'The ID of the cell to edit. When inserting a new cell, the new cell will be inserted after the cell with this ID, or at the beginning if not specified.',
+      ),
+    new_source: z.string().describe('The new source for the cell'),
+    cell_type: z
+      .enum(['code', 'markdown'])
+      .optional()
+      .describe(
+        'The type of the cell (code or markdown). If not specified, it defaults to the current cell type. If using edit_mode=insert, this is required.',
+      ),
+    edit_mode: z
+      .enum(['replace', 'insert', 'delete'])
+      .optional()
+      .describe(
+        'The type of edit to make (replace, insert, delete). Defaults to replace.',
+      ),
+  }),
+)
+type InputSchema = ReturnType<typeof inputSchema>
+
+export const outputSchema = lazySchema(() =>
+  z.object({
+    new_source: z
+      .string()
+      .describe('The new source code that was written to the cell'),
+    cell_id: z
+      .string()
+      .optional()
+      .describe('The ID of the cell that was edited'),
+    cell_type: z.enum(['code', 'markdown']).describe('The type of the cell'),
+    language: z.string().describe('The programming language of the notebook'),
+    edit_mode: z.string().describe('The edit mode that was used'),
+    error: z
+      .string()
+      .optional()
+      .describe('Error message if the operation failed'),
+    // Fields for attribution tracking
+    notebook_path: z.string().describe('The path to the notebook file'),
+    original_file: z
+      .string()
+      .describe('The original notebook content before modification'),
+    updated_file: z
+      .string()
+      .describe('The updated notebook content after modification'),
+  }),
+)
+type OutputSchema = ReturnType<typeof outputSchema>
+
+export type Output = z.infer<OutputSchema>
+
+export const NotebookEditTool = buildTool({
+  name: NOTEBOOK_EDIT_TOOL_NAME,
+  searchHint: 'edit Jupyter notebook cells (.ipynb)',
+  maxResultSizeChars: 100_000,
+  shouldDefer: true,
+  async description() {
+    return DESCRIPTION
+  },
+  async prompt() {
+    return PROMPT
+  },
+  userFacingName() {
+    return 'Edit Notebook'
+  },
+  getToolUseSummary,
+  getActivityDescription(input) {
+    const summary = getToolUseSummary(input)
+    return summary ? `Editing notebook ${summary}` : 'Editing notebook'
+  },
+  get inputSchema(): InputSchema {
+    return inputSchema()
+  },
+  get outputSchema(): OutputSchema {
+    return outputSchema()
+  },
+  toAutoClassifierInput(input) {
+    if (feature('TRANSCRIPT_CLASSIFIER')) {
+      const mode = input.edit_mode ?? 'replace'
+      return `${input.notebook_path} ${mode}: ${input.new_source}`
+    }
+    return ''
+  },
+  getPath(input): string {
+    return input.notebook_path
+  },
+  async checkPermissions(input, context): Promise<PermissionDecision> {
+    const appState = context.getAppState()
+    return checkWritePermissionForTool(
+      NotebookEditTool,
+      input,
+      appState.toolPermissionContext,
+    )
+  },
+  mapToolResultToToolResultBlockParam(
+    { cell_id, edit_mode, new_source, error },
+    toolUseID,
+  ) {
+    if (error) {
+      return {
+        tool_use_id: toolUseID,
+        type: 'tool_result',
+        content: error,
+        is_error: true,
+      }
+    }
+    switch (edit_mode) {
+      case 'replace':
+        return {
+          tool_use_id: toolUseID,
+          type: 'tool_result',
+          content: `Updated cell ${cell_id} with ${new_source}`,
+        }
+      case 'insert':
+        return {
+          tool_use_id: toolUseID,
+          type: 'tool_result',
+          content: `Inserted cell ${cell_id} with ${new_source}`,
+        }
+      case 'delete':
+        return {
+          tool_use_id: toolUseID,
+          type: 'tool_result',
+          content: `Deleted cell ${cell_id}`,
+        }
+      default:
+        return {
+          tool_use_id: toolUseID,
+          type: 'tool_result',
+          content: 'Unknown edit mode',
+        }
+    }
+  },
+  renderToolUseMessage,
+  renderToolUseRejectedMessage,
+  renderToolUseErrorMessage,
+  renderToolResultMessage,
+  async validateInput(
+    { notebook_path, cell_type, cell_id, edit_mode = 'replace' },
+    toolUseContext: ToolUseContext,
+  ) {
+    const fullPath = isAbsolute(notebook_path)
+      ? notebook_path
+      : resolve(getCwd(), notebook_path)
+
+    // SECURITY: Skip filesystem operations for UNC paths to prevent NTLM credential leaks.
+    if (fullPath.startsWith('\\\\') || fullPath.startsWith('//')) {
+      return { result: true }
+    }
+
+    if (extname(fullPath) !== '.ipynb') {
+      return {
+        result: false,
+        message:
+          'File must be a Jupyter notebook (.ipynb file). For editing other file types, use the FileEdit tool.',
+        errorCode: 2,
+      }
+    }
+
+    if (
+      edit_mode !== 'replace' &&
+      edit_mode !== 'insert' &&
+      edit_mode !== 'delete'
+    ) {
+      return {
+        result: false,
+        message: 'Edit mode must be replace, insert, or delete.',
+        errorCode: 4,
+      }
+    }
+
+    if (edit_mode === 'insert' && !cell_type) {
+      return {
+        result: false,
+        message: 'Cell type is required when using edit_mode=insert.',
+        errorCode: 5,
+      }
+    }
+
+    // Require Read-before-Edit (matches FileEditTool/FileWriteTool). Without
+    // this, the model could edit a notebook it never saw, or edit against a
+    // stale view after an external change — silent data loss.
+    const readTimestamp = toolUseContext.readFileState.get(fullPath)
+    if (!readTimestamp) {
+      return {
+        result: false,
+        message:
+          'File has not been read yet. Read it first before writing to it.',
+        errorCode: 9,
+      }
+    }
+    if (getFileModificationTime(fullPath) > readTimestamp.timestamp) {
+      return {
+        result: false,
+        message:
+          'File has been modified since read, either by the user or by a linter. Read it again before attempting to write it.',
+        errorCode: 10,
+      }
+    }
+
+    let content: string
+    try {
+      content = readFileSyncWithMetadata(fullPath).content
+    } catch (e) {
+      if (isENOENT(e)) {
+        return {
+          result: false,
+          message: 'Notebook file does not exist.',
+          errorCode: 1,
+        }
+      }
+      throw e
+    }
+    const notebook = safeParseJSON(content) as NotebookContent | null
+    if (!notebook) {
+      return {
+        result: false,
+        message: 'Notebook is not valid JSON.',
+        errorCode: 6,
+      }
+    }
+    if (!cell_id) {
+      if (edit_mode !== 'insert') {
+        return {
+          result: false,
+          message: 'Cell ID must be specified when not inserting a new cell.',
+          errorCode: 7,
+        }
+      }
+    } else {
+      // First try to find the cell by its actual ID
+      const cellIndex = notebook.cells.findIndex(cell => cell.id === cell_id)
+
+      if (cellIndex === -1) {
+        // If not found, try to parse as a numeric index (cell-N format)
+        const parsedCellIndex = parseCellId(cell_id)
+        if (parsedCellIndex !== undefined) {
+          if (!notebook.cells[parsedCellIndex]) {
+            return {
+              result: false,
+              message: `Cell with index ${parsedCellIndex} does not exist in notebook.`,
+              errorCode: 7,
+            }
+          }
+        } else {
+          return {
+            result: false,
+            message: `Cell with ID "${cell_id}" not found in notebook.`,
+            errorCode: 8,
+          }
+        }
+      }
+    }
+
+    return { result: true }
+  },
+  async call(
+    {
+      notebook_path,
+      new_source,
+      cell_id,
+      cell_type,
+      edit_mode: originalEditMode,
+    },
+    { readFileState, updateFileHistoryState },
+    _,
+    parentMessage,
+  ) {
+    const fullPath = isAbsolute(notebook_path)
+      ? notebook_path
+      : resolve(getCwd(), notebook_path)
+
+    if (fileHistoryEnabled()) {
+      await fileHistoryTrackEdit(
+        updateFileHistoryState,
+        fullPath,
+        parentMessage.uuid,
+      )
+    }
+
+    try {
+      // readFileSyncWithMetadata gives content + encoding + line endings in
+      // one safeResolvePath + readFileSync pass, replacing the previous
+      // detectFileEncoding + readFile + detectLineEndings chain (each of
+      // which redid safeResolvePath and/or a 4KB readSync).
+      const { content, encoding, lineEndings } =
+        readFileSyncWithMetadata(fullPath)
+      // Must use non-memoized jsonParse here: safeParseJSON caches by content
+      // string and returns a shared object reference, but we mutate the
+      // notebook in place below (cells.splice, targetCell.source = ...).
+      // Using the memoized version poisons the cache for validateInput() and
+      // any subsequent call() with the same file content.
+      let notebook: NotebookContent
+      try {
+        notebook = jsonParse(content) as NotebookContent
+      } catch {
+        return {
+          data: {
+            new_source,
+            cell_type: cell_type ?? 'code',
+            language: 'python',
+            edit_mode: 'replace',
+            error: 'Notebook is not valid JSON.',
+            cell_id,
+            notebook_path: fullPath,
+            original_file: '',
+            updated_file: '',
+          },
+        }
+      }
+
+      let cellIndex
+      if (!cell_id) {
+        cellIndex = 0 // Default to inserting at the beginning if no cell_id is provided
+      } else {
+        // First try to find the cell by its actual ID
+        cellIndex = notebook.cells.findIndex(cell => cell.id === cell_id)
+
+        // If not found, try to parse as a numeric index (cell-N format)
+        if (cellIndex === -1) {
+          const parsedCellIndex = parseCellId(cell_id)
+          if (parsedCellIndex !== undefined) {
+            cellIndex = parsedCellIndex
+          }
+        }
+
+        if (originalEditMode === 'insert') {
+          cellIndex += 1 // Insert after the cell with this ID
+        }
+      }
+
+      // Convert replace to insert if trying to replace one past the end
+      let edit_mode = originalEditMode
+      if (edit_mode === 'replace' && cellIndex === notebook.cells.length) {
+        edit_mode = 'insert'
+        if (!cell_type) {
+          cell_type = 'code' // Default to code if no cell_type specified
+        }
+      }
+
+      const language = notebook.metadata.language_info?.name ?? 'python'
+      let new_cell_id = undefined
+      if (
+        notebook.nbformat > 4 ||
+        (notebook.nbformat === 4 && notebook.nbformat_minor >= 5)
+      ) {
+        if (edit_mode === 'insert') {
+          new_cell_id = Math.random().toString(36).substring(2, 15)
+        } else if (cell_id !== null) {
+          new_cell_id = cell_id
+        }
+      }
+
+      if (edit_mode === 'delete') {
+        // Delete the specified cell
+        notebook.cells.splice(cellIndex, 1)
+      } else if (edit_mode === 'insert') {
+        let new_cell: NotebookCell
+        if (cell_type === 'markdown') {
+          new_cell = {
+            cell_type: 'markdown',
+            id: new_cell_id,
+            source: new_source,
+            metadata: {},
+          }
+        } else {
+          new_cell = {
+            cell_type: 'code',
+            id: new_cell_id,
+            source: new_source,
+            metadata: {},
+            execution_count: null,
+            outputs: [],
+          }
+        }
+        // Insert the new cell
+        notebook.cells.splice(cellIndex, 0, new_cell)
+      } else {
+        // Find the specified cell
+        const targetCell = notebook.cells[cellIndex]! // validateInput ensures cell_number is in bounds
+        targetCell.source = new_source
+        if (targetCell.cell_type === 'code') {
+          // Reset execution count and clear outputs since cell was modified
+          targetCell.execution_count = null
+          targetCell.outputs = []
+        }
+        if (cell_type && cell_type !== targetCell.cell_type) {
+          targetCell.cell_type = cell_type
+        }
+      }
+      // Write back to file
+      const IPYNB_INDENT = 1
+      const updatedContent = jsonStringify(notebook, null, IPYNB_INDENT)
+      writeTextContent(fullPath, updatedContent, encoding, lineEndings)
+      // Update readFileState with post-write mtime (matches FileEditTool/
+      // FileWriteTool). offset:undefined breaks FileReadTool's dedup match —
+      // without this, Read→NotebookEdit→Read in the same millisecond would
+      // return the file_unchanged stub against stale in-context content.
+      readFileState.set(fullPath, {
+        content: updatedContent,
+        timestamp: getFileModificationTime(fullPath),
+        offset: undefined,
+        limit: undefined,
+      })
+      const data = {
+        new_source,
+        cell_type: cell_type ?? 'code',
+        language,
+        edit_mode: edit_mode ?? 'replace',
+        cell_id: new_cell_id || undefined,
+        error: '',
+        notebook_path: fullPath,
+        original_file: content,
+        updated_file: updatedContent,
+      }
+      return {
+        data,
+      }
+    } catch (error) {
+      if (error instanceof Error) {
+        const data = {
+          new_source,
+          cell_type: cell_type ?? 'code',
+          language: 'python',
+          edit_mode: 'replace',
+          error: error.message,
+          cell_id,
+          notebook_path: fullPath,
+          original_file: '',
+          updated_file: '',
+        }
+        return {
+          data,
+        }
+      }
+      const data = {
+        new_source,
+        cell_type: cell_type ?? 'code',
+        language: 'python',
+        edit_mode: 'replace',
+        error: 'Unknown error occurred while editing notebook',
+        cell_id,
+        notebook_path: fullPath,
+        original_file: '',
+        updated_file: '',
+      }
+      return {
+        data,
+      }
+    }
+  },
+} satisfies ToolDef<InputSchema, Output>)