From 2603ec0b59857bb9329223f18705241dd615bbef Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Fri, 20 Jun 2025 10:50:30 +0200 Subject: [PATCH 1/4] fix: wrap ERROR and WARN logging in vim.schedule to prevent fast event context errors (#54) --- lua/claudecode/logger.lua | 8 +- tests/selection_test.lua | 3 + tests/unit/logger_spec.lua | 175 +++++++++++++++++++++++++++++++++++++ 3 files changed, 184 insertions(+), 2 deletions(-) create mode 100644 tests/unit/logger_spec.lua diff --git a/lua/claudecode/logger.lua b/lua/claudecode/logger.lua index 44418a3..1a8969d 100644 --- a/lua/claudecode/logger.lua +++ b/lua/claudecode/logger.lua @@ -69,9 +69,13 @@ local function log(level, component, message_parts) end if level == M.levels.ERROR then - vim.notify(prefix .. " " .. message, vim.log.levels.ERROR, { title = "ClaudeCode Error" }) + vim.schedule(function() + vim.notify(prefix .. " " .. message, vim.log.levels.ERROR, { title = "ClaudeCode Error" }) + end) elseif level == M.levels.WARN then - vim.notify(prefix .. " " .. message, vim.log.levels.WARN, { title = "ClaudeCode Warning" }) + vim.schedule(function() + vim.notify(prefix .. " " .. message, vim.log.levels.WARN, { title = "ClaudeCode Warning" }) + end) else -- For INFO, DEBUG, TRACE, use nvim_echo to avoid flooding notifications, -- to make them appear in :messages, and wrap in vim.schedule diff --git a/tests/selection_test.lua b/tests/selection_test.lua index 1ecf0f6..ee59383 100644 --- a/tests/selection_test.lua +++ b/tests/selection_test.lua @@ -3,6 +3,9 @@ if not _G.vim then schedule_wrap = function(fn) return fn end, + schedule = function(fn) + fn() + end, _buffers = {}, _windows = {}, _commands = {}, diff --git a/tests/unit/logger_spec.lua b/tests/unit/logger_spec.lua new file mode 100644 index 0000000..572617c --- /dev/null +++ b/tests/unit/logger_spec.lua @@ -0,0 +1,175 @@ +-- luacheck: globals expect +require("tests.busted_setup") + +describe("Logger", function() + local logger + local original_vim_schedule + local original_vim_notify + local original_nvim_echo + local scheduled_calls = {} + local notify_calls = {} + local echo_calls = {} + + local function setup() + package.loaded["claudecode.logger"] = nil + + -- Mock vim.schedule to track calls + original_vim_schedule = vim.schedule + vim.schedule = function(fn) + table.insert(scheduled_calls, fn) + -- Immediately execute the function for testing + fn() + end + + -- Mock vim.notify to track calls + original_vim_notify = vim.notify + vim.notify = function(msg, level, opts) + table.insert(notify_calls, { msg = msg, level = level, opts = opts }) + end + + -- Mock nvim_echo to track calls + original_nvim_echo = vim.api.nvim_echo + vim.api.nvim_echo = function(chunks, history, opts) + table.insert(echo_calls, { chunks = chunks, history = history, opts = opts }) + end + + logger = require("claudecode.logger") + + -- Set log level to TRACE to enable all logging levels for testing + logger.setup({ log_level = "trace" }) + end + + local function teardown() + vim.schedule = original_vim_schedule + vim.notify = original_vim_notify + vim.api.nvim_echo = original_nvim_echo + scheduled_calls = {} + notify_calls = {} + echo_calls = {} + end + + before_each(function() + setup() + end) + + after_each(function() + teardown() + end) + + describe("error logging", function() + it("should wrap error calls in vim.schedule", function() + logger.error("test", "error message") + + -- Should have made one scheduled call + expect(#scheduled_calls).to_be(1) + + -- Should have called vim.notify with error level + expect(#notify_calls).to_be(1) + expect(notify_calls[1].level).to_be(vim.log.levels.ERROR) + assert_contains(notify_calls[1].msg, "error message") + end) + + it("should handle error calls without component", function() + logger.error("error message") + + expect(#scheduled_calls).to_be(1) + expect(#notify_calls).to_be(1) + assert_contains(notify_calls[1].msg, "error message") + end) + end) + + describe("warn logging", function() + it("should wrap warn calls in vim.schedule", function() + logger.warn("test", "warning message") + + -- Should have made one scheduled call + expect(#scheduled_calls).to_be(1) + + -- Should have called vim.notify with warn level + expect(#notify_calls).to_be(1) + expect(notify_calls[1].level).to_be(vim.log.levels.WARN) + assert_contains(notify_calls[1].msg, "warning message") + end) + + it("should handle warn calls without component", function() + logger.warn("warning message") + + expect(#scheduled_calls).to_be(1) + expect(#notify_calls).to_be(1) + assert_contains(notify_calls[1].msg, "warning message") + end) + end) + + describe("info logging", function() + it("should wrap info calls in vim.schedule", function() + logger.info("test", "info message") + + -- Should have made one scheduled call + expect(#scheduled_calls).to_be(1) + + -- Should have called nvim_echo instead of notify + expect(#echo_calls).to_be(1) + expect(#notify_calls).to_be(0) + assert_contains(echo_calls[1].chunks[1][1], "info message") + end) + end) + + describe("debug logging", function() + it("should wrap debug calls in vim.schedule", function() + logger.debug("test", "debug message") + + -- Should have made one scheduled call + expect(#scheduled_calls).to_be(1) + + -- Should have called nvim_echo instead of notify + expect(#echo_calls).to_be(1) + expect(#notify_calls).to_be(0) + assert_contains(echo_calls[1].chunks[1][1], "debug message") + end) + end) + + describe("trace logging", function() + it("should wrap trace calls in vim.schedule", function() + logger.trace("test", "trace message") + + -- Should have made one scheduled call + expect(#scheduled_calls).to_be(1) + + -- Should have called nvim_echo instead of notify + expect(#echo_calls).to_be(1) + expect(#notify_calls).to_be(0) + assert_contains(echo_calls[1].chunks[1][1], "trace message") + end) + end) + + describe("fast event context safety", function() + it("should not call vim API functions directly", function() + -- Simulate a fast event context by removing the mocked functions + -- and ensuring no direct calls are made + local direct_notify_called = false + local direct_echo_called = false + + vim.notify = function() + direct_notify_called = true + end + + vim.api.nvim_echo = function() + direct_echo_called = true + end + + vim.schedule = function(fn) + -- Don't execute the function, just verify it was scheduled + table.insert(scheduled_calls, fn) + end + + logger.error("test", "error in fast context") + logger.warn("test", "warn in fast context") + logger.info("test", "info in fast context") + + -- All should be scheduled, none should be called directly + expect(#scheduled_calls).to_be(3) + expect(direct_notify_called).to_be_false() + expect(direct_echo_called).to_be_false() + end) + end) +end) From 3f09b5163af9ce6e1869ea3220c77f8229d934c5 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Fri, 20 Jun 2025 14:22:04 +0200 Subject: [PATCH 2/4] docs: update PROTOCOL.md with complete VS Code tool specs and streamline README (#55) --- PROTOCOL.md | 371 ++++++++++++++++++++++++++++++++++++++---- README.md | 428 ++++++++----------------------------------------- dev-config.lua | 1 + 3 files changed, 408 insertions(+), 392 deletions(-) diff --git a/PROTOCOL.md b/PROTOCOL.md index 384ccf9..6614a08 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -132,46 +132,361 @@ According to the MCP spec, Claude should be able to call tools, but **current im ## Available MCP Tools -The extensions register these tools that Claude can (theoretically) call: +The VS Code extension registers 12 tools that Claude can call. Here's the complete specification: -### Core Tools +### 1. openFile -1. **openFile** - Open a file and optionally select text +**Description**: Open a file in the editor and optionally select a range of text - ```json - { - "filePath": "/path/to/file.js", - "startText": "function hello", // Find and select from this text - "endText": "}" // To this text - } - ``` +**Input**: -2. **openDiff** - Show a diff and wait for user action (blocking!) +```json +{ + "filePath": "/path/to/file.js", + "preview": false, + "startText": "function hello", + "endText": "}", + "selectToEndOfLine": false, + "makeFrontmost": true +} +``` + +- `filePath` (string, required): Path to the file to open +- `preview` (boolean, default: false): Whether to open in preview mode +- `startText` (string, optional): Text pattern to find selection start +- `endText` (string, optional): Text pattern to find selection end +- `selectToEndOfLine` (boolean, default: false): Extend selection to end of line +- `makeFrontmost` (boolean, default: true): Make the file the active editor tab + +**Output**: When `makeFrontmost=true`, returns simple message: + +```json +{ + "content": [ + { + "type": "text", + "text": "Opened file: /path/to/file.js" + } + ] +} +``` + +When `makeFrontmost=false`, returns detailed JSON: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": true, \"filePath\": \"/absolute/path/to/file.js\", \"languageId\": \"javascript\", \"lineCount\": 42}" + } + ] +} +``` + +### 2. openDiff + +**Description**: Open a git diff for the file (blocking operation) + +**Input**: + +```json +{ + "old_file_path": "/path/to/original.js", + "new_file_path": "/path/to/modified.js", + "new_file_contents": "// Modified content...", + "tab_name": "Proposed changes" +} +``` + +- `old_file_path` (string): Path to original file +- `new_file_path` (string): Path to new file +- `new_file_contents` (string): Contents of the new file +- `tab_name` (string): Tab name for the diff view + +**Output**: Returns MCP-formatted response: + +```json +{ + "content": [ + { + "type": "text", + "text": "FILE_SAVED" + } + ] +} +``` + +or + +```json +{ + "content": [ + { + "type": "text", + "text": "DIFF_REJECTED" + } + ] +} +``` + +Based on whether the user saves or rejects the diff. + +### 3. getCurrentSelection - ```json - { - "old_file_path": "/path/to/original.js", - "new_file_path": "/path/to/modified.js", - "new_file_contents": "// Modified content...", - "tab_name": "Proposed changes" - } - ``` +**Description**: Get the current text selection in the active editor + +**Input**: None + +**Output**: Returns JSON-stringified selection data: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": true, \"text\": \"selected content\", \"filePath\": \"/path/to/file\", \"selection\": {\"start\": {\"line\": 0, \"character\": 0}, \"end\": {\"line\": 0, \"character\": 10}}}" + } + ] +} +``` + +Or when no active editor: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": false, \"message\": \"No active editor found\"}" + } + ] +} +``` + +### 4. getLatestSelection + +**Description**: Get the most recent text selection (even if not in active editor) + +**Input**: None + +**Output**: JSON-stringified selection data or `{success: false, message: "No selection available"}` + +### 5. getOpenEditors + +**Description**: Get information about currently open editors + +**Input**: None + +**Output**: Returns JSON-stringified array of open tabs: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"tabs\": [{\"uri\": \"file:///path/to/file\", \"isActive\": true, \"label\": \"filename.ext\", \"languageId\": \"javascript\", \"isDirty\": false}]}" + } + ] +} +``` + +### 6. getWorkspaceFolders + +**Description**: Get all workspace folders currently open in the IDE + +**Input**: None + +**Output**: Returns JSON-stringified workspace information: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": true, \"folders\": [{\"name\": \"project-name\", \"uri\": \"file:///path/to/workspace\", \"path\": \"/path/to/workspace\"}], \"rootPath\": \"/path/to/workspace\"}" + } + ] +} +``` + +### 7. getDiagnostics + +**Description**: Get language diagnostics from VS Code + +**Input**: + +```json +{ + "uri": "file:///path/to/file.js" +} +``` + +- `uri` (string, optional): File URI to get diagnostics for. If not provided, gets diagnostics for all files. + +**Output**: Returns JSON-stringified array of diagnostics per file: + +```json +{ + "content": [ + { + "type": "text", + "text": "[{\"uri\": \"file:///path/to/file\", \"diagnostics\": [{\"message\": \"Error message\", \"severity\": \"Error\", \"range\": {\"start\": {\"line\": 0, \"character\": 0}}, \"source\": \"typescript\"}]}]" + } + ] +} +``` + +### 8. checkDocumentDirty + +**Description**: Check if a document has unsaved changes (is dirty) + +**Input**: + +```json +{ + "filePath": "/path/to/file.js" +} +``` + +- `filePath` (string, required): Path to the file to check + +**Output**: Returns document dirty status: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": true, \"filePath\": \"/path/to/file.js\", \"isDirty\": true, \"isUntitled\": false}" + } + ] +} +``` + +Or when document not open: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": false, \"message\": \"Document not open: /path/to/file.js\"}" + } + ] +} +``` + +### 9. saveDocument + +**Description**: Save a document with unsaved changes + +**Input**: + +```json +{ + "filePath": "/path/to/file.js" +} +``` + +- `filePath` (string, required): Path to the file to save + +**Output**: Returns save operation result: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": true, \"filePath\": \"/path/to/file.js\", \"saved\": true, \"message\": \"Document saved successfully\"}" + } + ] +} +``` + +Or when document not open: + +```json +{ + "content": [ + { + "type": "text", + "text": "{\"success\": false, \"message\": \"Document not open: /path/to/file.js\"}" + } + ] +} +``` + +### 10. close_tab + +**Description**: Close a tab by name + +**Input**: + +```json +{ + "tab_name": "filename.js" +} +``` + +- `tab_name` (string, required): Name of the tab to close + +**Output**: Returns `{content: [{type: "text", text: "TAB_CLOSED"}]}` + +### 11. closeAllDiffTabs + +**Description**: Close all diff tabs in the editor + +**Input**: None + +**Output**: Returns `{content: [{type: "text", text: "CLOSED_${count}_DIFF_TABS"}]}` + +### 12. executeCode + +**Description**: Execute Python code in the Jupyter kernel for the current notebook file + +**Input**: + +```json +{ + "code": "print('Hello, World!')" +} +``` + +- `code` (string, required): The code to be executed on the kernel + +**Output**: Returns execution results with mixed content types: + +```json +{ + "content": [ + { + "type": "text", + "text": "Hello, World!" + }, + { + "type": "image", + "data": "base64_encoded_image_data", + "mimeType": "image/png" + } + ] +} +``` - Returns `FILE_SAVED` or `DIFF_REJECTED` based on user action. +**Notes**: -3. **getCurrentSelection** - Get the current text selection -4. **getOpenEditors** - List all open files -5. **getWorkspaceFolders** - Get project folders -6. **getDiagnostics** - Get errors/warnings from the IDE -7. **saveDocument** - Save a file -8. **close_tab** - Close a tab by name (note the inconsistent naming!) +- All code executed will persist across calls unless the kernel is restarted +- Avoid declaring variables or modifying kernel state unless explicitly requested +- Only available when working with Jupyter notebooks +- Can return multiple content types including text output and images ### Implementation Notes - Most tools follow camelCase naming except `close_tab` (uses snake_case) -- The `openDiff` tool is unique - it's **blocking** and waits for user interaction +- The `openDiff` tool is **blocking** and waits for user interaction - Tools return MCP-formatted responses with content arrays -- There's also `executeCode` for Jupyter notebooks in the VS Code extension +- All schemas use Zod validation in the VS Code extension +- Selection-related tools work with the current editor state ## Building Your Own Integration diff --git a/README.md b/README.md index 968a8fe..358ebe9 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,11 @@ ![Neovim version](https://img.shields.io/badge/Neovim-0.8%2B-green) ![Status](https://img.shields.io/badge/Status-beta-blue) +> ⚠️ **Important**: IDE integrations are currently broken in Claude Code releases newer than v1.0.27. Please use [Claude Code v1.0.27](https://www.npmjs.com/package/@anthropic-ai/claude-code/v/1.0.27) or older until these issues are resolved: +> +> - [Claude Code not detecting IDE integrations #2299](https://github.com/anthropics/claude-code/issues/2299) +> - [IDE integration broken after update #2295](https://github.com/anthropics/claude-code/issues/2295) + **The first Neovim IDE integration for Claude Code** — bringing Anthropic's AI coding assistant to your favorite editor with a pure Lua implementation. > 🎯 **TL;DR:** When Anthropic released Claude Code with VS Code and JetBrains support, I reverse-engineered their extension and built this Neovim plugin. This plugin implements the same WebSocket-based MCP protocol, giving Neovim users the same AI-powered coding experience. @@ -20,33 +25,12 @@ When Anthropic released Claude Code, they only supported VS Code and JetBrains. - ⚡ **First to Market** — Beat Anthropic to releasing Neovim support - 🛠️ **Built with AI** — Used Claude to reverse-engineer Claude's own protocol -## Quick Demo - -```vim -" Launch Claude Code in a split -:ClaudeCode - -" Claude now sees your current file and selections in real-time! - -" Send visual selection as context -:'<,'>ClaudeCodeSend - -" Claude can open files, show diffs, and more -``` - -## Requirements - -- Neovim >= 0.8.0 -- [Claude Code CLI](https://docs.anthropic.com/en/docs/claude-code) installed -- Optional: [folke/snacks.nvim](https://github.com/folke/snacks.nvim) for enhanced terminal support - ## Installation -Using [lazy.nvim](https://github.com/folke/lazy.nvim): - ```lua { "coder/claudecode.nvim", + dependencies = { "folke/snacks.nvim" }, config = true, keys = { { "a", nil, desc = "AI/Claude Code" }, @@ -54,6 +38,7 @@ Using [lazy.nvim](https://github.com/folke/lazy.nvim): { "af", "ClaudeCodeFocus", desc = "Focus Claude" }, { "ar", "ClaudeCode --resume", desc = "Resume Claude" }, { "aC", "ClaudeCode --continue", desc = "Continue Claude" }, + { "ab", "ClaudeCodeAdd %", desc = "Add current buffer" }, { "as", "ClaudeCodeSend", mode = "v", desc = "Send to Claude" }, { "as", @@ -68,162 +53,71 @@ Using [lazy.nvim](https://github.com/folke/lazy.nvim): } ``` -That's it! For more configuration options, see [Advanced Setup](#advanced-setup). +That's it! The plugin will auto-configure everything else. + +## Requirements + +- Neovim >= 0.8.0 +- [Claude Code CLI](https://docs.anthropic.com/en/docs/claude-code) installed +- [folke/snacks.nvim](https://github.com/folke/snacks.nvim) for enhanced terminal support + +## Quick Demo + +```vim +" Launch Claude Code in a split +:ClaudeCode + +" Claude now sees your current file and selections in real-time! + +" Send visual selection as context +:'<,'>ClaudeCodeSend + +" Claude can open files, show diffs, and more +``` ## Usage 1. **Launch Claude**: Run `:ClaudeCode` to open Claude in a split terminal 2. **Send context**: - Select text in visual mode and use `as` to send it to Claude - - In `nvim-tree` or `neo-tree`, press `as` on a file to add it to Claude's context + - In `nvim-tree`/`neo-tree`/`oil.nvim`, press `as` on a file to add it to Claude's context 3. **Let Claude work**: Claude can now: - See your current file and selections in real-time - Open files in your editor - Show diffs with proposed changes - Access diagnostics and workspace info -## Commands - -- `:ClaudeCode [arguments]` - Toggle the Claude Code terminal window (simple show/hide behavior) -- `:ClaudeCodeFocus [arguments]` - Smart focus/toggle Claude terminal (switches to terminal if not focused, hides if focused) -- `:ClaudeCode --resume` - Resume a previous Claude conversation -- `:ClaudeCode --continue` - Continue Claude conversation -- `:ClaudeCodeSend` - Send current visual selection to Claude, or add files from tree explorer -- `:ClaudeCodeTreeAdd` - Add selected file(s) from tree explorer to Claude context (also available via ClaudeCodeSend) -- `:ClaudeCodeAdd [start-line] [end-line]` - Add a specific file or directory to Claude context by path with optional line range -- `:ClaudeCodeDiffAccept` - Accept the current diff changes (equivalent to `aa`) -- `:ClaudeCodeDiffDeny` - Deny/reject the current diff changes (equivalent to `ad`) - -### Toggle Behavior - -- **`:ClaudeCode`** - Simple toggle: Always show/hide terminal regardless of current focus -- **`:ClaudeCodeFocus`** - Smart focus: Focus terminal if not active, hide if currently focused - -### Tree Integration - -The `as` keybinding has context-aware behavior: +## Key Commands -- **In normal buffers (visual mode)**: Sends selected text to Claude -- **In nvim-tree/neo-tree/oil.nvim buffers**: Adds the file under cursor (or selected files) to Claude's context - -This allows you to quickly add entire files to Claude's context for review, refactoring, or discussion. - -#### Features - -- **Single file**: Place cursor on any file and press `as` -- **Multiple files**: Select multiple files (using tree plugin's selection features or visual selection in oil.nvim) and press `as` -- **Smart detection**: Automatically detects whether you're in nvim-tree, neo-tree, or oil.nvim -- **Error handling**: Clear feedback if no files are selected or if tree plugins aren't available - -### Direct File Addition - -The `:ClaudeCodeAdd` command allows you to add files or directories directly by path, with optional line range specification: - -```vim -:ClaudeCodeAdd src/main.lua -:ClaudeCodeAdd ~/projects/myproject/ -:ClaudeCodeAdd ./README.md -:ClaudeCodeAdd src/main.lua 50 100 " Lines 50-100 only -:ClaudeCodeAdd config.lua 25 " Only line 25 -``` - -#### Features - -- **Path completion**: Tab completion for file and directory paths -- **Path expansion**: Supports `~` for home directory and relative paths -- **Line range support**: Optionally specify start and end lines for files (ignored for directories) -- **Validation**: Checks that files and directories exist before adding, validates line numbers -- **Flexible**: Works with both individual files and entire directories +- `:ClaudeCode` - Toggle the Claude Code terminal window +- `:ClaudeCodeFocus` - Smart focus/toggle Claude terminal +- `:ClaudeCodeSend` - Send current visual selection to Claude +- `:ClaudeCodeAdd [start-line] [end-line]` - Add specific file to Claude context with optional line range +- `:ClaudeCodeDiffAccept` - Accept diff changes +- `:ClaudeCodeDiffDeny` - Reject diff changes ## Working with Diffs -When Claude proposes changes to your files, the plugin opens a native Neovim diff view showing the original file alongside the proposed changes. You have several options to accept or reject these changes: - -### Accepting Changes - -- **`:w` (save)** - Accept the changes and apply them to your file -- **`aa`** - Accept the changes using the dedicated keymap (configured in LazyVim spec) - -You can edit the proposed changes in the right-hand diff buffer before accepting them. This allows you to modify Claude's suggestions or make additional tweaks before applying the final version to your file. - -Both methods signal Claude Code to apply the changes to your file, after which the plugin automatically reloads the affected buffers to show the updated content. +When Claude proposes changes, the plugin opens a native Neovim diff view: -### Rejecting Changes +- **Accept**: `:w` (save) or `aa` +- **Reject**: `:q` or `ad` -- **`:q` or `:close`** - Close the diff view to reject the changes -- **`ad`** - Reject changes using the dedicated keymap (configured in LazyVim spec) -- **`:bdelete` or `:bwipeout`** - Delete the diff buffer to reject changes - -When you reject changes, the diff view closes and the original file remains unchanged. - -### Accepting/Rejecting from Claude Code Terminal - -You can also navigate to the Claude Code terminal window and accept or reject diffs directly from within Claude's interface. This provides an alternative way to manage diffs without using the Neovim-specific keymaps. - -### Customizing Diff Keymaps - -The diff keymaps are configured in the LazyVim spec and can be customized by modifying the `keys` table: - -```lua -{ - "coder/claudecode.nvim", - config = true, - keys = { - -- ... other keymaps ... - - -- Customize diff keymaps to avoid conflicts (e.g., with debugger) - { "ya", "ClaudeCodeDiffAccept", desc = "Accept diff" }, - { "yn", "ClaudeCodeDiffDeny", desc = "Deny diff" }, - - -- Or disable them entirely by omitting them from the keys table - }, -} -``` - -The commands `ClaudeCodeDiffAccept` and `ClaudeCodeDiffDeny` work only in diff buffers created by the plugin and will show a warning if used elsewhere. - -### How It Works - -The plugin uses a signal-based approach where accepting or rejecting a diff sends a message to Claude Code rather than directly modifying files. This ensures consistency and allows Claude Code to handle the actual file operations while the plugin manages the user interface and buffer reloading. - -#### Examples - -```vim -" Add entire files -:ClaudeCodeAdd src/components/Header.tsx -:ClaudeCodeAdd ~/.config/nvim/init.lua - -" Add entire directories (line numbers ignored) -:ClaudeCodeAdd tests/ -:ClaudeCodeAdd ../other-project/ - -" Add specific line ranges -:ClaudeCodeAdd src/main.lua 50 100 " Lines 50 through 100 -:ClaudeCodeAdd config.lua 25 " Only line 25 -:ClaudeCodeAdd utils.py 1 50 " First 50 lines -:ClaudeCodeAdd README.md 10 20 " Just lines 10-20 - -" Path expansion works with line ranges -:ClaudeCodeAdd ~/project/src/app.js 100 200 -:ClaudeCodeAdd ./relative/path.lua 30 -``` +You can edit Claude's suggestions before accepting them. ## How It Works This plugin creates a WebSocket server that Claude Code CLI connects to, implementing the same protocol as the official VS Code extension. When you launch Claude, it automatically detects Neovim and gains full access to your editor. -### The Protocol - -The extensions use a WebSocket-based variant of the Model Context Protocol (MCP) that only Claude Code supports. The plugin: +The protocol uses a WebSocket-based variant of MCP (Model Context Protocol) that: 1. Creates a WebSocket server on a random port 2. Writes a lock file to `~/.claude/ide/[port].lock` with connection info 3. Sets environment variables that tell Claude where to connect 4. Implements MCP tools that Claude can call -For the full technical details and protocol documentation, see [PROTOCOL.md](./PROTOCOL.md). - 📖 **[Read the full reverse-engineering story →](./STORY.md)** +🔧 **[Complete protocol documentation →](./PROTOCOL.md)** ## Architecture @@ -237,44 +131,7 @@ Built with pure Lua and zero external dependencies: For deep technical details, see [ARCHITECTURE.md](./ARCHITECTURE.md). -## Contributing - -See [DEVELOPMENT.md](./DEVELOPMENT.md) for build instructions and development guidelines. Tests can be run with `make test`. - -## Configuration - -### Quick Setup - -For most users, the default configuration is sufficient: - -```lua -{ - "coder/claudecode.nvim", - dependencies = { - "folke/snacks.nvim", -- optional - }, - config = true, - keys = { - { "a", nil, desc = "AI/Claude Code" }, - { "ac", "ClaudeCode", desc = "Toggle Claude" }, - { "af", "ClaudeCodeFocus", desc = "Focus Claude" }, - { "ar", "ClaudeCode --resume", desc = "Resume Claude" }, - { "aC", "ClaudeCode --continue", desc = "Continue Claude" }, - { "as", "ClaudeCodeSend", mode = "v", desc = "Send to Claude" }, - { - "as", - "ClaudeCodeTreeAdd", - desc = "Add file", - ft = { "NvimTree", "neo-tree", "oil" }, - }, - -- Diff management - { "aa", "ClaudeCodeDiffAccept", desc = "Accept diff" }, - { "ad", "ClaudeCodeDiffDeny", desc = "Deny diff" }, - }, -} -``` - -### Advanced Configuration +## Advanced Configuration
Complete configuration options @@ -282,208 +139,51 @@ For most users, the default configuration is sufficient: ```lua { "coder/claudecode.nvim", - dependencies = { - "folke/snacks.nvim", -- Optional for enhanced terminal - }, - keys = { - { "a", nil, desc = "AI/Claude Code" }, - { "ac", "ClaudeCode", desc = "Toggle Claude" }, - { "af", "ClaudeCodeFocus", desc = "Focus Claude" }, - { "ar", "ClaudeCode --resume", desc = "Resume Claude" }, - { "aC", "ClaudeCode --continue", desc = "Continue Claude" }, - { "as", "ClaudeCodeSend", mode = "v", desc = "Send to Claude" }, - { - "as", - "ClaudeCodeTreeAdd", - desc = "Add file", - ft = { "NvimTree", "neo-tree", "oil" }, - }, - -- Diff management - { "aa", "ClaudeCodeDiffAccept", desc = "Accept diff" }, - { "ad", "ClaudeCodeDiffDeny", desc = "Deny diff" }, - }, + dependencies = { "folke/snacks.nvim" }, opts = { -- Server Configuration - port_range = { min = 10000, max = 65535 }, -- WebSocket server port range - auto_start = true, -- Auto-start server on Neovim startup - log_level = "info", -- "trace", "debug", "info", "warn", "error" - terminal_cmd = nil, -- Custom terminal command (default: "claude") + port_range = { min = 10000, max = 65535 }, + auto_start = true, + log_level = "info", -- "trace", "debug", "info", "warn", "error" + terminal_cmd = nil, -- Custom terminal command (default: "claude") -- Selection Tracking - track_selection = true, -- Enable real-time selection tracking - visual_demotion_delay_ms = 50, -- Delay before demoting visual selection (ms) - - -- Connection Management - connection_wait_delay = 200, -- Wait time after connection before sending queued @ mentions (ms) - connection_timeout = 10000, -- Max time to wait for Claude Code connection (ms) - queue_timeout = 5000, -- Max time to keep @ mentions in queue (ms) + track_selection = true, + visual_demotion_delay_ms = 50, -- Terminal Configuration terminal = { - split_side = "right", -- "left" or "right" - split_width_percentage = 0.30, -- Width as percentage (0.0 to 1.0) - provider = "auto", -- "auto", "snacks", or "native" - show_native_term_exit_tip = true, -- Show exit tip for native terminal - auto_close = true, -- Auto-close terminal after command completion + split_side = "right", -- "left" or "right" + split_width_percentage = 0.30, + provider = "auto", -- "auto", "snacks", or "native" + auto_close = true, }, -- Diff Integration diff_opts = { - auto_close_on_accept = true, -- Close diff view after accepting changes - show_diff_stats = true, -- Show diff statistics - vertical_split = true, -- Use vertical split for diffs - open_in_current_tab = true, -- Open diffs in current tab vs new tab - }, - }, -} -``` - -
- -### Configuration Options Explained - -#### Server Options - -- **`port_range`**: Port range for the WebSocket server that Claude connects to -- **`auto_start`**: Whether to automatically start the integration when Neovim starts -- **`terminal_cmd`**: Override the default "claude" command (useful for custom Claude installations) -- **`log_level`**: Controls verbosity of plugin logs - -#### Selection Tracking - -- **`track_selection`**: Enables real-time selection updates sent to Claude -- **`visual_demotion_delay_ms`**: Time to wait before switching from visual selection to cursor position tracking - -#### Connection Management - -- **`connection_wait_delay`**: Prevents overwhelming Claude with rapid @ mentions after connection -- **`connection_timeout`**: How long to wait for Claude to connect before giving up -- **`queue_timeout`**: How long to keep queued @ mentions before discarding them - -#### Terminal Configuration - -- **`split_side`**: Which side to open the terminal split (`"left"` or `"right"`) -- **`split_width_percentage`**: Terminal width as a fraction of screen width (0.1 = 10%, 0.5 = 50%) -- **`provider`**: Terminal implementation to use: - - `"auto"`: Try snacks.nvim, fallback to native - - `"snacks"`: Force snacks.nvim (requires folke/snacks.nvim) - - `"native"`: Use built-in Neovim terminal -- **`show_native_term_exit_tip`**: Show help text for exiting native terminal -- **`auto_close`**: Automatically close terminal when commands finish - -#### Diff Options - -- **`auto_close_on_accept`**: Close diff view after accepting changes with `:w` or `aa` -- **`show_diff_stats`**: Display diff statistics (lines added/removed) -- **`vertical_split`**: Use vertical split layout for diffs -- **`open_in_current_tab`**: Open diffs in current tab instead of creating new tabs - -### Example Configurations - -#### Minimal Configuration - -```lua -{ - "coder/claudecode.nvim", - keys = { - { "a", nil, desc = "AI/Claude Code" }, - { "ac", "ClaudeCode", desc = "Toggle Claude" }, - { "af", "ClaudeCodeFocus", desc = "Focus Claude" }, - { "ar", "ClaudeCode --resume", desc = "Resume Claude" }, - { "aC", "ClaudeCode --continue", desc = "Continue Claude" }, - { "as", "ClaudeCodeSend", mode = "v", desc = "Send to Claude" }, - { - "as", - "ClaudeCodeTreeAdd", - desc = "Add file", - ft = { "NvimTree", "neo-tree", "oil" }, + auto_close_on_accept = true, + vertical_split = true, + open_in_current_tab = true, }, - -- Diff management - { "aa", "ClaudeCodeDiffAccept", desc = "Accept diff" }, - { "ad", "ClaudeCodeDiffDeny", desc = "Deny diff" }, - }, - opts = { - log_level = "warn", -- Reduce log verbosity - auto_start = false, -- Manual startup only }, -} -``` - -#### Power User Configuration - -```lua -{ - "coder/claudecode.nvim", keys = { - { "a", nil, desc = "AI/Claude Code" }, - { "ac", "ClaudeCode", desc = "Toggle Claude" }, - { "af", "ClaudeCodeFocus", desc = "Focus Claude" }, - { "ar", "ClaudeCode --resume", desc = "Resume Claude" }, - { "aC", "ClaudeCode --continue", desc = "Continue Claude" }, - { "as", "ClaudeCodeSend", mode = "v", desc = "Send to Claude" }, - { - "as", - "ClaudeCodeTreeAdd", - desc = "Add file", - ft = { "NvimTree", "neo-tree", "oil" }, - }, - -- Diff management - { "aa", "ClaudeCodeDiffAccept", desc = "Accept diff" }, - { "ad", "ClaudeCodeDiffDeny", desc = "Deny diff" }, - }, - opts = { - log_level = "debug", - visual_demotion_delay_ms = 100, -- Slower selection demotion - connection_wait_delay = 500, -- Longer delay for @ mention batching - terminal = { - split_side = "left", - split_width_percentage = 0.4, -- Wider terminal - provider = "snacks", - auto_close = false, -- Keep terminal open to review output - }, - diff_opts = { - vertical_split = false, -- Horizontal diffs - open_in_current_tab = false, -- New tabs for diffs - }, + -- Your keymaps here }, } ``` -#### Custom Claude Installation - -```lua -{ - "coder/claudecode.nvim", - keys = { - { "a", nil, desc = "AI/Claude Code" }, - { "ac", "ClaudeCode", desc = "Toggle Claude" }, - { "af", "ClaudeCodeFocus", desc = "Focus Claude" }, - { "ar", "ClaudeCode --resume", desc = "Resume Claude" }, - { "aC", "ClaudeCode --continue", desc = "Continue Claude" }, - { "as", "ClaudeCodeSend", mode = "v", desc = "Send to Claude" }, - { - "as", - "ClaudeCodeTreeAdd", - desc = "Add file", - ft = { "NvimTree", "neo-tree", "oil" }, - }, - -- Diff management - { "aa", "ClaudeCodeDiffAccept", desc = "Accept diff" }, - { "ad", "ClaudeCodeDiffDeny", desc = "Deny diff" }, - }, - opts = { - terminal_cmd = "/opt/claude/bin/claude", -- Custom Claude path - port_range = { min = 20000, max = 25000 }, -- Different port range - }, -} -``` + ## Troubleshooting - **Claude not connecting?** Check `:ClaudeCodeStatus` and verify lock file exists in `~/.claude/ide/` -- **Need debug logs?** Set `log_level = "debug"` in setup +- **Need debug logs?** Set `log_level = "debug"` in opts - **Terminal issues?** Try `provider = "native"` if using snacks.nvim +## Contributing + +See [DEVELOPMENT.md](./DEVELOPMENT.md) for build instructions and development guidelines. Tests can be run with `make test`. + ## License [MIT](LICENSE) diff --git a/dev-config.lua b/dev-config.lua index 5e6dde6..ecb3489 100644 --- a/dev-config.lua +++ b/dev-config.lua @@ -17,6 +17,7 @@ return { { "aC", "ClaudeCode --continue", desc = "Continue Claude" }, -- Context sending + { "ab", "ClaudeCodeAdd %", desc = "Add current buffer" }, { "as", "ClaudeCodeSend", mode = "v", desc = "Send to Claude" }, { "as", From 3fe2c1b4497f9edfe9c5608a3b7e5bae87273244 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Fri, 20 Jun 2025 18:08:18 +0200 Subject: [PATCH 3/4] feat: implement WebSocket authentication system with UUID tokens (#56) --- CLAUDE.md | 151 ++++++++++++++- PROTOCOL.md | 23 ++- lua/claudecode/init.lua | 61 +++++- lua/claudecode/lockfile.lua | 118 ++++++++++- lua/claudecode/server/client.lua | 48 ++++- lua/claudecode/server/handshake.lua | 39 +++- lua/claudecode/server/init.lua | 26 ++- lua/claudecode/server/tcp.lua | 7 +- scripts/claude_interactive.sh | 5 +- scripts/lib_claude.sh | 65 ++++++- scripts/lib_ws_persistent.sh | 11 +- tests/integration/mcp_tools_spec.lua | 113 +++++++++++ tests/lockfile_test.lua | 128 +++++++++++- tests/unit/init_spec.lua | 5 +- tests/unit/server/handshake_spec.lua | 279 +++++++++++++++++++++++++++ tests/unit/server_spec.lua | 136 +++++++++++++ 16 files changed, 1177 insertions(+), 38 deletions(-) create mode 100644 tests/unit/server/handshake_spec.lua diff --git a/CLAUDE.md b/CLAUDE.md index a47aa1a..abd634c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -22,10 +22,16 @@ claudecode.nvim - A Neovim plugin that implements the same WebSocket-based MCP p ### Build Commands +- `make` - **RECOMMENDED**: Run formatting, linting, and testing (complete validation) - `make all` - Run check and format (default target) +- `make test` - Run all tests using busted with coverage +- `make check` - Check Lua syntax and run luacheck +- `make format` - Format code with stylua (or nix fmt if available) - `make clean` - Remove generated test files - `make help` - Show available commands +**Best Practice**: Always use `make` at the end of editing sessions for complete validation. + ### Development with Nix - `nix develop` - Enter development shell with all dependencies @@ -45,11 +51,20 @@ claudecode.nvim - A Neovim plugin that implements the same WebSocket-based MCP p ### WebSocket Server Implementation - **TCP Server**: `server/tcp.lua` handles port binding and connections -- **Handshake**: `server/handshake.lua` processes HTTP upgrade requests +- **Handshake**: `server/handshake.lua` processes HTTP upgrade requests with authentication - **Frame Processing**: `server/frame.lua` implements RFC 6455 WebSocket frames - **Client Management**: `server/client.lua` manages individual connections - **Utils**: `server/utils.lua` provides base64, SHA-1, XOR operations in pure Lua +#### Authentication System + +The WebSocket server implements secure authentication using: + +- **UUID v4 Tokens**: Generated per session with enhanced entropy +- **Header-based Auth**: Uses `x-claude-code-ide-authorization` header +- **Lock File Discovery**: Tokens stored in `~/.claude/ide/[port].lock` for Claude CLI +- **MCP Compliance**: Follows official Claude Code IDE authentication protocol + ### MCP Tools Architecture Tools are registered with JSON schemas and handlers. MCP-exposed tools include: @@ -76,14 +91,125 @@ Tests are organized in three layers: Test files follow the pattern `*_spec.lua` or `*_test.lua` and use the busted framework. +### Test Organization Principles + +- **Isolation**: Each test should be independent and not rely on external state +- **Mocking**: Use comprehensive mocking for vim APIs and external dependencies +- **Coverage**: Aim for both positive and negative test cases, edge cases included +- **Performance**: Tests should run quickly to encourage frequent execution +- **Clarity**: Test names should clearly describe what behavior is being verified + +## Authentication Testing + +The plugin implements authentication using UUID v4 tokens that are generated for each server session and stored in lock files. This ensures secure connections between Claude CLI and the Neovim WebSocket server. + +### Testing Authentication Features + +**Lock File Authentication Tests** (`tests/lockfile_test.lua`): + +- Auth token generation and uniqueness validation +- Lock file creation with authentication tokens +- Reading auth tokens from existing lock files +- Error handling for missing or invalid tokens + +**WebSocket Handshake Authentication Tests** (`tests/unit/server/handshake_spec.lua`): + +- Valid authentication token acceptance +- Invalid/missing token rejection +- Edge cases (empty tokens, malformed headers, length limits) +- Case-insensitive header handling + +**Server Integration Tests** (`tests/unit/server_spec.lua`): + +- Server startup with authentication tokens +- Auth token state management during server lifecycle +- Token validation throughout server operations + +**End-to-End Authentication Tests** (`tests/integration/mcp_tools_spec.lua`): + +- Complete authentication flow from server start to tool execution +- Authentication state persistence across operations +- Concurrent operations with authentication enabled + +### Manual Authentication Testing + +**Test Script Authentication Support**: + +```bash +# Test scripts automatically detect and use authentication tokens +cd scripts/ +./claude_interactive.sh # Automatically reads auth token from lock file +``` + +**Authentication Flow Testing**: + +1. Start the plugin: `:ClaudeCodeStart` +2. Check lock file contains `authToken`: `cat ~/.claude/ide/*.lock | jq .authToken` +3. Test WebSocket connection with auth: Use test scripts in `scripts/` directory +4. Verify authentication in logs: Set `log_level = "debug"` in config + +**Testing Authentication Failures**: + +```bash +# Test invalid auth token (should fail) +websocat ws://localhost:PORT --header "x-claude-code-ide-authorization: invalid-token" + +# Test missing auth header (should fail) +websocat ws://localhost:PORT + +# Test valid auth token (should succeed) +websocat ws://localhost:PORT --header "x-claude-code-ide-authorization: $(cat ~/.claude/ide/*.lock | jq -r .authToken)" +``` + +### Authentication Logging + +Enable detailed authentication logging by setting: + +```lua +require("claudecode").setup({ + log_level = "debug" -- Shows auth token generation, validation, and failures +}) +``` + +Log levels for authentication events: + +- **DEBUG**: Server startup authentication state, client connections, handshake processing, auth token details +- **WARN**: Authentication failures during handshake +- **ERROR**: Auth token generation failures, handshake response errors + +### Logging Best Practices + +- **Connection Events**: Use DEBUG level for routine connection establishment/teardown +- **Authentication Flow**: Use DEBUG for successful auth, WARN for failures +- **User-Facing Events**: Use INFO sparingly for events users need to know about +- **System Errors**: Use ERROR for failures that require user attention + ## Development Notes +### Technical Requirements + - Plugin requires Neovim >= 0.8.0 - Uses only Neovim built-ins for WebSocket implementation (vim.loop, vim.json, vim.schedule) -- Lock files are created at `~/.claude/ide/[port].lock` for Claude CLI discovery -- WebSocket server only accepts local connections for security +- Zero external dependencies for core functionality + +### Security Considerations + +- WebSocket server only accepts local connections (127.0.0.1) for security +- Authentication tokens are UUID v4 with enhanced entropy +- Lock files created at `~/.claude/ide/[port].lock` for Claude CLI discovery +- All authentication events are logged for security auditing + +### Performance Optimizations + - Selection tracking is debounced to reduce overhead +- WebSocket frame processing optimized for JSON-RPC payload sizes +- Connection pooling and cleanup to prevent resource leaks + +### Integration Support + - Terminal integration supports both snacks.nvim and native Neovim terminal +- Compatible with popular file explorers (nvim-tree, oil.nvim) +- Visual selection tracking across different selection modes ## Release Process @@ -134,6 +260,23 @@ make rg "0\.1\.0" . # Should only show CHANGELOG.md historical entries ``` -## CRITICAL: Pre-commit Requirements +## Development Workflow + +### Pre-commit Requirements **ALWAYS run `make` before committing any changes.** This runs code quality checks and formatting that must pass for CI to succeed. Never skip this step - many PRs fail CI because contributors don't run the build commands before committing. + +### Recommended Development Flow + +1. **Start Development**: Use existing tests and documentation to understand the system +2. **Make Changes**: Follow existing patterns and conventions in the codebase +3. **Validate Work**: Run `make` to ensure formatting, linting, and tests pass +4. **Document Changes**: Update relevant documentation (this file, PROTOCOL.md, etc.) +5. **Commit**: Only commit after successful `make` execution + +### Code Quality Standards + +- **Test Coverage**: Maintain comprehensive test coverage (currently 314+ tests) +- **Zero Warnings**: All code must pass luacheck with 0 warnings/errors +- **Consistent Formatting**: Use `nix fmt` or `stylua` for consistent code style +- **Documentation**: Update CLAUDE.md for architectural changes, PROTOCOL.md for protocol changes diff --git a/PROTOCOL.md b/PROTOCOL.md index 6614a08..042579c 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -23,7 +23,8 @@ The IDE writes a discovery file to `~/.claude/ide/[port].lock`: "pid": 12345, // IDE process ID "workspaceFolders": ["/path/to/project"], // Open folders "ideName": "VS Code", // or "Neovim", "IntelliJ", etc. - "transport": "ws" // WebSocket transport + "transport": "ws", // WebSocket transport + "authToken": "550e8400-e29b-41d4-a716-446655440000" // Random UUID for authentication } ``` @@ -38,6 +39,16 @@ When launching Claude, the IDE sets: Claude reads the lock files, finds the matching port from the environment, and connects to the WebSocket server. +## Authentication + +When Claude connects to the IDE's WebSocket server, it must authenticate using the token from the lock file. The authentication happens via a custom WebSocket header: + +``` +x-claude-code-ide-authorization: 550e8400-e29b-41d4-a716-446655440000 +``` + +The IDE validates this header against the `authToken` value from the lock file. If the token doesn't match, the connection is rejected. + ## The Protocol Communication uses WebSocket with JSON-RPC 2.0 messages: @@ -503,11 +514,13 @@ local server = create_websocket_server("127.0.0.1", random_port) ```lua -- ~/.claude/ide/[port].lock +local auth_token = generate_uuid() -- Generate random UUID local lock_data = { pid = vim.fn.getpid(), workspaceFolders = { vim.fn.getcwd() }, ideName = "YourEditor", - transport = "ws" + transport = "ws", + authToken = auth_token } write_json(lock_path, lock_data) ``` @@ -523,6 +536,12 @@ claude # Claude will now connect! ### 4. Handle Messages ```lua +-- Validate authentication on WebSocket handshake +function validate_auth(headers) + local auth_header = headers["x-claude-code-ide-authorization"] + return auth_header == auth_token +end + -- Send selection updates send_message({ jsonrpc = "2.0", diff --git a/lua/claudecode/init.lua b/lua/claudecode/init.lua index e318f28..f673899 100644 --- a/lua/claudecode/init.lua +++ b/lua/claudecode/init.lua @@ -67,6 +67,7 @@ local default_config = { --- @field config ClaudeCode.Config The current plugin configuration. --- @field server table|nil The WebSocket server instance. --- @field port number|nil The port the server is running on. +--- @field auth_token string|nil The authentication token for the current session. --- @field initialized boolean Whether the plugin has been initialized. --- @field queued_mentions table[] Array of queued @ mentions waiting for connection. --- @field connection_timer table|nil Timer for connection timeout. @@ -76,6 +77,7 @@ M.state = { config = vim.deepcopy(default_config), server = nil, port = nil, + auth_token = nil, initialized = false, queued_mentions = {}, connection_timer = nil, @@ -357,26 +359,70 @@ function M.start(show_startup_notification) end local server = require("claudecode.server.init") - local success, result = server.start(M.state.config) + local lockfile = require("claudecode.lockfile") + + -- Generate auth token first so we can pass it to the server + local auth_token + local auth_success, auth_result = pcall(function() + return lockfile.generate_auth_token() + end) + + if not auth_success then + local error_msg = "Failed to generate authentication token: " .. (auth_result or "unknown error") + logger.error("init", error_msg) + return false, error_msg + end + + auth_token = auth_result + + -- Validate the generated auth token + if not auth_token or type(auth_token) ~= "string" or #auth_token < 10 then + local error_msg = "Invalid authentication token generated" + logger.error("init", error_msg) + return false, error_msg + end + + local success, result = server.start(M.state.config, auth_token) if not success then - logger.error("init", "Failed to start Claude Code integration: " .. result) - return false, result + local error_msg = "Failed to start Claude Code server: " .. (result or "unknown error") + if result and result:find("auth") then + error_msg = error_msg .. " (authentication related)" + end + logger.error("init", error_msg) + return false, error_msg end M.state.server = server M.state.port = tonumber(result) + M.state.auth_token = auth_token - local lockfile = require("claudecode.lockfile") - local lock_success, lock_result = lockfile.create(M.state.port) + local lock_success, lock_result, returned_auth_token = lockfile.create(M.state.port, auth_token) if not lock_success then server.stop() M.state.server = nil M.state.port = nil + M.state.auth_token = nil + + local error_msg = "Failed to create lock file: " .. (lock_result or "unknown error") + if lock_result and lock_result:find("auth") then + error_msg = error_msg .. " (authentication token issue)" + end + logger.error("init", error_msg) + return false, error_msg + end - logger.error("init", "Failed to create lock file: " .. lock_result) - return false, lock_result + -- Verify that the auth token in the lock file matches what we generated + if returned_auth_token ~= auth_token then + server.stop() + M.state.server = nil + M.state.port = nil + M.state.auth_token = nil + + local error_msg = "Authentication token mismatch between server and lock file" + logger.error("init", error_msg) + return false, error_msg end if M.state.config.track_selection then @@ -422,6 +468,7 @@ function M.stop() M.state.server = nil M.state.port = nil + M.state.auth_token = nil -- Clear any queued @ mentions when server stops clear_mention_queue() diff --git a/lua/claudecode/lockfile.lua b/lua/claudecode/lockfile.lua index 12792a9..1590dad 100644 --- a/lua/claudecode/lockfile.lua +++ b/lua/claudecode/lockfile.lua @@ -9,15 +9,68 @@ local M = {} --- Path to the lock file directory M.lock_dir = vim.fn.expand("~/.claude/ide") +-- Track if random seed has been initialized +local random_initialized = false + +--- Generate a random UUID for authentication +---@return string uuid A randomly generated UUID string +local function generate_auth_token() + -- Initialize random seed only once + if not random_initialized then + local seed = os.time() + vim.fn.getpid() + -- Add more entropy if available + if vim.loop and vim.loop.hrtime then + seed = seed + (vim.loop.hrtime() % 1000000) + end + math.randomseed(seed) + + -- Call math.random a few times to "warm up" the generator + for _ = 1, 10 do + math.random() + end + random_initialized = true + end + + -- Generate UUID v4 format: xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx + local template = "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx" + local uuid = template:gsub("[xy]", function(c) + local v = (c == "x") and math.random(0, 15) or math.random(8, 11) + return string.format("%x", v) + end) + + -- Validate generated UUID format + if not uuid:match("^[0-9a-f]+-[0-9a-f]+-[0-9a-f]+-[0-9a-f]+-[0-9a-f]+$") then + error("Generated invalid UUID format: " .. uuid) + end + + if #uuid ~= 36 then + error("Generated UUID has invalid length: " .. #uuid .. " (expected 36)") + end + + return uuid +end + +--- Generate a new authentication token +---@return string auth_token A newly generated authentication token +function M.generate_auth_token() + return generate_auth_token() +end + --- Create the lock file for a specified WebSocket port ---@param port number The port number for the WebSocket server +---@param auth_token string|nil Optional pre-generated auth token (generates new one if not provided) ---@return boolean success Whether the operation was successful ---@return string result_or_error The lock file path if successful, or error message if failed -function M.create(port) +---@return string? auth_token The authentication token if successful +function M.create(port, auth_token) if not port or type(port) ~= "number" then return false, "Invalid port number" end + if port < 1 or port > 65535 then + return false, "Port number out of valid range (1-65535): " .. tostring(port) + end + local ok, err = pcall(function() return vim.fn.mkdir(M.lock_dir, "p") end) @@ -29,6 +82,24 @@ function M.create(port) local lock_path = M.lock_dir .. "/" .. port .. ".lock" local workspace_folders = M.get_workspace_folders() + if not auth_token then + local auth_success, auth_result = pcall(generate_auth_token) + if not auth_success then + return false, "Failed to generate authentication token: " .. (auth_result or "unknown error") + end + auth_token = auth_result + else + -- Validate provided auth_token + if type(auth_token) ~= "string" then + return false, "Authentication token must be a string, got " .. type(auth_token) + end + if #auth_token < 10 then + return false, "Authentication token too short (minimum 10 characters)" + end + if #auth_token > 500 then + return false, "Authentication token too long (maximum 500 characters)" + end + end -- Prepare lock file content local lock_content = { @@ -36,6 +107,7 @@ function M.create(port) workspaceFolders = workspace_folders, ideName = "Neovim", transport = "ws", + authToken = auth_token, } local json @@ -65,7 +137,7 @@ function M.create(port) return false, "Failed to write lock file: " .. (write_err or "unknown error") end - return true, lock_path + return true, lock_path, auth_token end --- Remove the lock file for the given port @@ -98,6 +170,7 @@ end ---@param port number The port number of the WebSocket server ---@return boolean success Whether the operation was successful ---@return string result_or_error The lock file path if successful, or error message if failed +---@return string? auth_token The authentication token if successful function M.update(port) if not port or type(port) ~= "number" then return false, "Invalid port number" @@ -114,6 +187,47 @@ function M.update(port) return M.create(port) end +--- Read the authentication token from a lock file +---@param port number The port number of the WebSocket server +---@return boolean success Whether the operation was successful +---@return string? auth_token The authentication token if successful, or nil if failed +---@return string? error Error message if operation failed +function M.get_auth_token(port) + if not port or type(port) ~= "number" then + return false, nil, "Invalid port number" + end + + local lock_path = M.lock_dir .. "/" .. port .. ".lock" + + if vim.fn.filereadable(lock_path) == 0 then + return false, nil, "Lock file does not exist: " .. lock_path + end + + local file = io.open(lock_path, "r") + if not file then + return false, nil, "Failed to open lock file: " .. lock_path + end + + local content = file:read("*all") + file:close() + + if not content or content == "" then + return false, nil, "Lock file is empty: " .. lock_path + end + + local ok, lock_data = pcall(vim.json.decode, content) + if not ok or type(lock_data) ~= "table" then + return false, nil, "Failed to parse lock file JSON: " .. lock_path + end + + local auth_token = lock_data.authToken + if not auth_token or type(auth_token) ~= "string" then + return false, nil, "No valid auth token found in lock file" + end + + return true, auth_token, nil +end + --- Get active LSP clients using available API ---@return table Array of LSP clients local function get_lsp_clients() diff --git a/lua/claudecode/server/client.lua b/lua/claudecode/server/client.lua index ebe3bfc..031cd2a 100644 --- a/lua/claudecode/server/client.lua +++ b/lua/claudecode/server/client.lua @@ -1,6 +1,7 @@ ---@brief WebSocket client connection management local frame = require("claudecode.server.frame") local handshake = require("claudecode.server.handshake") +local logger = require("claudecode.logger") local M = {} @@ -38,16 +39,55 @@ end ---@param on_message function Callback for complete messages: function(client, message_text) ---@param on_close function Callback for client close: function(client, code, reason) ---@param on_error function Callback for errors: function(client, error_msg) -function M.process_data(client, data, on_message, on_close, on_error) +---@param auth_token string|nil Expected authentication token for validation +function M.process_data(client, data, on_message, on_close, on_error, auth_token) client.buffer = client.buffer .. data if not client.handshake_complete then local complete, request, remaining = handshake.extract_http_request(client.buffer) if complete then - local success, response_from_handshake, _ = handshake.process_handshake(request) + logger.debug("client", "Processing WebSocket handshake for client:", client.id) + + -- Log if auth token is required + if auth_token then + logger.debug("client", "Authentication required for client:", client.id) + else + logger.debug("client", "No authentication required for client:", client.id) + end + + local success, response_from_handshake, _ = handshake.process_handshake(request, auth_token) + + -- Log authentication results + if success then + if auth_token then + logger.debug("client", "Client authenticated successfully:", client.id) + else + logger.debug("client", "Client handshake completed (no auth required):", client.id) + end + else + -- Log specific authentication failure details + if auth_token and response_from_handshake:find("auth") then + logger.warn( + "client", + "Authentication failed for client " + .. client.id + .. ": " + .. (response_from_handshake:match("Bad WebSocket upgrade request: (.+)") or "unknown auth error") + ) + else + logger.warn( + "client", + "WebSocket handshake failed for client " + .. client.id + .. ": " + .. (response_from_handshake:match("HTTP/1.1 %d+ (.+)") or "unknown handshake error") + ) + end + end client.tcp_handle:write(response_from_handshake, function(err) if err then + logger.error("client", "Failed to send handshake response to client " .. client.id .. ": " .. err) on_error(client, "Failed to send handshake response: " .. err) return end @@ -56,12 +96,14 @@ function M.process_data(client, data, on_message, on_close, on_error) client.handshake_complete = true client.state = "connected" client.buffer = remaining + logger.debug("client", "WebSocket connection established for client:", client.id) if #client.buffer > 0 then - M.process_data(client, "", on_message, on_close, on_error) + M.process_data(client, "", on_message, on_close, on_error, auth_token) end else client.state = "closing" + logger.debug("client", "Closing connection for client due to failed handshake:", client.id) vim.schedule(function() client.tcp_handle:close() end) diff --git a/lua/claudecode/server/handshake.lua b/lua/claudecode/server/handshake.lua index a7ec162..4f04f3d 100644 --- a/lua/claudecode/server/handshake.lua +++ b/lua/claudecode/server/handshake.lua @@ -5,9 +5,10 @@ local M = {} ---@brief Check if an HTTP request is a valid WebSocket upgrade request ---@param request string The HTTP request string +---@param expected_auth_token string|nil Expected authentication token for validation ---@return boolean valid True if it's a valid WebSocket upgrade request ---@return table|string headers_or_error Headers table if valid, error message if not -function M.validate_upgrade_request(request) +function M.validate_upgrade_request(request, expected_auth_token) local headers = utils.parse_http_headers(request) -- Check for required headers @@ -33,6 +34,37 @@ function M.validate_upgrade_request(request) return false, "Invalid Sec-WebSocket-Key format" end + -- Validate authentication token if required + if expected_auth_token then + -- Check if expected_auth_token is valid + if type(expected_auth_token) ~= "string" or expected_auth_token == "" then + return false, "Server configuration error: invalid expected authentication token" + end + + local auth_header = headers["x-claude-code-ide-authorization"] + if not auth_header then + return false, "Missing authentication header: x-claude-code-ide-authorization" + end + + -- Check for empty auth header + if auth_header == "" then + return false, "Authentication token too short (min 10 characters)" + end + + -- Check for suspicious auth header lengths + if #auth_header > 500 then + return false, "Authentication token too long (max 500 characters)" + end + + if #auth_header < 10 then + return false, "Authentication token too short (min 10 characters)" + end + + if auth_header ~= expected_auth_token then + return false, "Invalid authentication token" + end + end + return true, headers end @@ -131,10 +163,11 @@ end ---@brief Process a complete WebSocket handshake ---@param request string The HTTP request string +---@param expected_auth_token string|nil Expected authentication token for validation ---@return boolean success True if handshake was successful ---@return string response The HTTP response to send ---@return table|nil headers The parsed headers if successful -function M.process_handshake(request) +function M.process_handshake(request, expected_auth_token) -- Check if it's a valid WebSocket endpoint request if not M.is_websocket_endpoint(request) then local response = M.create_error_response(404, "WebSocket endpoint not found") @@ -142,7 +175,7 @@ function M.process_handshake(request) end -- Validate the upgrade request - local is_valid_upgrade, validation_payload = M.validate_upgrade_request(request) ---@type boolean, table|string + local is_valid_upgrade, validation_payload = M.validate_upgrade_request(request, expected_auth_token) ---@type boolean, table|string if not is_valid_upgrade then assert(type(validation_payload) == "string", "validation_payload should be a string on error") local error_message = validation_payload diff --git a/lua/claudecode/server/init.lua b/lua/claudecode/server/init.lua index 0d764ef..b8762ae 100644 --- a/lua/claudecode/server/init.lua +++ b/lua/claudecode/server/init.lua @@ -11,12 +11,14 @@ local M = {} ---@class ServerState ---@field server table|nil The TCP server instance ---@field port number|nil The port server is running on +---@field auth_token string|nil The authentication token for validating connections ---@field clients table A list of connected clients ---@field handlers table Message handlers by method name ---@field ping_timer table|nil Timer for sending pings M.state = { server = nil, port = nil, + auth_token = nil, clients = {}, handlers = {}, ping_timer = nil, @@ -24,13 +26,24 @@ M.state = { ---@brief Initialize the WebSocket server ---@param config table Configuration options +---@param auth_token string|nil The authentication token for validating connections ---@return boolean success Whether server started successfully ---@return number|string port_or_error Port number or error message -function M.start(config) +function M.start(config, auth_token) if M.state.server then return false, "Server already running" end + M.state.auth_token = auth_token + + -- Log authentication state + if auth_token then + logger.debug("server", "Starting WebSocket server with authentication enabled") + logger.debug("server", "Auth token length:", #auth_token) + else + logger.debug("server", "Starting WebSocket server WITHOUT authentication (insecure)") + end + M.register_handlers() tools.setup(M) @@ -41,7 +54,13 @@ function M.start(config) end, on_connect = function(client) M.state.clients[client.id] = client - logger.debug("server", "WebSocket client connected:", client.id) + + -- Log connection with auth status + if M.state.auth_token then + logger.debug("server", "Authenticated WebSocket client connected:", client.id) + else + logger.debug("server", "WebSocket client connected (no auth):", client.id) + end -- Notify main module about new connection for queue processing local main_module = require("claudecode") @@ -68,7 +87,7 @@ function M.start(config) end, } - local server, error_msg = tcp_server.create_server(config, callbacks) + local server, error_msg = tcp_server.create_server(config, callbacks, M.state.auth_token) if not server then return false, error_msg or "Unknown server creation error" end @@ -104,6 +123,7 @@ function M.stop() M.state.server = nil M.state.port = nil + M.state.auth_token = nil M.state.clients = {} return true diff --git a/lua/claudecode/server/tcp.lua b/lua/claudecode/server/tcp.lua index ef3f30a..5b7462a 100644 --- a/lua/claudecode/server/tcp.lua +++ b/lua/claudecode/server/tcp.lua @@ -7,6 +7,7 @@ local M = {} ---@class TCPServer ---@field server table The vim.loop TCP server handle ---@field port number The port the server is listening on +---@field auth_token string|nil The authentication token for validating connections ---@field clients table Table of connected clients (client_id -> WebSocketClient) ---@field on_message function Callback for WebSocket messages ---@field on_connect function Callback for new connections @@ -47,9 +48,10 @@ end ---@brief Create and start a TCP server ---@param config table Server configuration ---@param callbacks table Callback functions +---@param auth_token string|nil Authentication token for validating connections ---@return TCPServer|nil server The server object, or nil on error ---@return string|nil error Error message if failed -function M.create_server(config, callbacks) +function M.create_server(config, callbacks, auth_token) local port = M.find_available_port(config.port_range.min, config.port_range.max) if not port then return nil, "No available ports in range " .. config.port_range.min .. "-" .. config.port_range.max @@ -64,6 +66,7 @@ function M.create_server(config, callbacks) local server = { server = tcp_server, port = port, + auth_token = auth_token, clients = {}, on_message = callbacks.on_message or function() end, on_connect = callbacks.on_connect or function() end, @@ -138,7 +141,7 @@ function M._handle_new_connection(server) end, function(cl, error_msg) server.on_error("Client " .. cl.id .. " error: " .. error_msg) M._remove_client(server, cl) - end) + end, server.auth_token) end) -- Notify about new connection diff --git a/scripts/claude_interactive.sh b/scripts/claude_interactive.sh index 1726179..8f7333a 100755 --- a/scripts/claude_interactive.sh +++ b/scripts/claude_interactive.sh @@ -28,13 +28,14 @@ if ! claude_is_running; then exit 1 fi -# Get WebSocket URL +# Get WebSocket URL and authentication info WS_URL=$(get_claude_ws_url) PORT=$(find_claude_lockfile) +AUTH_TOKEN=$(get_claude_auth_token "$PORT") # Initialize WebSocket connection echo -e "${BLUE}Initializing WebSocket connection to ${WS_URL}...${NC}" -if ! ws_connect "$WS_URL" "$CONN_ID"; then +if ! ws_connect "$WS_URL" "$CONN_ID" "$AUTH_TOKEN"; then echo -e "${RED}Failed to establish connection.${NC}" exit 1 fi diff --git a/scripts/lib_claude.sh b/scripts/lib_claude.sh index b31db0b..486474a 100755 --- a/scripts/lib_claude.sh +++ b/scripts/lib_claude.sh @@ -61,6 +61,48 @@ get_claude_ws_url() { echo "ws://localhost:$port" } +# Get the authentication token from a Claude Code lock file +# Usage: AUTH_TOKEN=$(get_claude_auth_token "$PORT") +get_claude_auth_token() { + local port="$1" + + if [[ -z $port ]]; then + echo >&2 "Error: Port number required" + return 1 + fi + + local lock_file="$CLAUDE_LOCKFILE_DIR/$port.lock" + + if [[ ! -f $lock_file ]]; then + echo >&2 "Error: Lock file not found: $lock_file" + return 1 + fi + + # Extract authToken from JSON using jq if available, otherwise basic parsing + if command -v jq >/dev/null 2>&1; then + local auth_token + auth_token=$(jq -r '.authToken // empty' "$lock_file" 2>/dev/null) + + if [[ -z $auth_token ]]; then + echo >&2 "Error: No authToken found in lock file" + return 1 + fi + + echo "$auth_token" + else + # Fallback parsing without jq + local auth_token + auth_token=$(grep -o '"authToken"[[:space:]]*:[[:space:]]*"[^"]*"' "$lock_file" | sed 's/.*"authToken"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/') + + if [[ -z $auth_token ]]; then + echo >&2 "Error: No authToken found in lock file (install jq for better JSON parsing)" + return 1 + fi + + echo "$auth_token" + fi +} + # Create a JSON-RPC request message (with ID) # Usage: MSG=$(create_message "method_name" '{"param":"value"}' "request-id") create_message() { @@ -123,10 +165,11 @@ create_init_message() { } # Send a message to the Claude Code WebSocket and get the response -# Usage: RESPONSE=$(send_claude_message "$MESSAGE" "$WS_URL") +# Usage: RESPONSE=$(send_claude_message "$MESSAGE" "$WS_URL" "$AUTH_TOKEN") send_claude_message() { local message="$1" local ws_url="${2:-}" + local auth_token="${3:-}" local timeout="${CLAUDE_WS_TIMEOUT:-5}" # Auto-detect WS URL if not provided @@ -134,8 +177,26 @@ send_claude_message() { ws_url=$(get_claude_ws_url) fi + # Auto-detect auth token if not provided + if [ -z "$auth_token" ]; then + local port + port=$(find_claude_lockfile) + if [[ $port =~ ^[0-9]+$ ]]; then + auth_token=$(get_claude_auth_token "$port") + fi + fi + + # Build websocat command with optional auth header + local websocat_cmd="websocat --protocol permessage-deflate --text" + + if [ -n "$auth_token" ]; then + websocat_cmd="$websocat_cmd --header 'x-claude-code-ide-authorization: $auth_token'" + fi + + websocat_cmd="$websocat_cmd '$ws_url' --no-close" + # Send message and get response with timeout - timeout "$timeout" bash -c "echo -n '$message' | websocat --protocol permessage-deflate --text '$ws_url' --no-close" 2>/dev/null || + timeout "$timeout" bash -c "echo -n '$message' | $websocat_cmd" 2>/dev/null || echo '{"error":{"code":-32000,"message":"Timeout waiting for response"}}' } diff --git a/scripts/lib_ws_persistent.sh b/scripts/lib_ws_persistent.sh index 51e6bce..dd92bbb 100755 --- a/scripts/lib_ws_persistent.sh +++ b/scripts/lib_ws_persistent.sh @@ -11,10 +11,11 @@ declare -A WS_CONNECTIONS declare -A WS_REQUEST_FILES # Start a persistent WebSocket connection -# ws_connect URL [CONN_ID] +# ws_connect URL [CONN_ID] [AUTH_TOKEN] ws_connect() { local url="$1" local conn_id="${2:-default}" + local auth_token="${3:-}" # Cleanup any existing connection with this ID ws_disconnect "$conn_id" @@ -41,7 +42,13 @@ ws_connect() { # 2. Writes all server responses to response_file ( # Note: The -E flag makes websocat exit when the file is closed - tail -f "$request_file" | websocat -t "$url" | tee -a "$response_file" >"$log_file" & + if [ -n "$auth_token" ]; then + # Use websocat with auth header - avoid eval by constructing command safely + tail -f "$request_file" | websocat -t --header "x-claude-code-ide-authorization: $auth_token" "$url" | tee -a "$response_file" >"$log_file" & + else + # Use websocat without auth header + tail -f "$request_file" | websocat -t "$url" | tee -a "$response_file" >"$log_file" & + fi # Save PID echo $! >"$pid_file" diff --git a/tests/integration/mcp_tools_spec.lua b/tests/integration/mcp_tools_spec.lua index 1c60e45..0b3ce90 100644 --- a/tests/integration/mcp_tools_spec.lua +++ b/tests/integration/mcp_tools_spec.lua @@ -709,5 +709,118 @@ describe("MCP Tools Integration", function() end) end) + describe("Authentication Flow Integration", function() + local test_auth_token = "550e8400-e29b-41d4-a716-446655440000" + local config = { + port_range = { + min = 10000, + max = 65535, + }, + } + + -- Ensure clean state before each test + before_each(function() + if server.state.server then + server.stop() + end + end) + + -- Clean up after each test + after_each(function() + if server.state.server then + server.stop() + end + end) + + it("should start server with auth token", function() + -- Start server with authentication + local success, port = server.start(config, test_auth_token) + expect(success).to_be_true() + expect(server.state.auth_token).to_be(test_auth_token) + expect(type(port)).to_be("number") + + -- Verify server is running with auth + local status = server.get_status() + expect(status.running).to_be_true() + expect(status.port).to_be(port) + + -- Clean up + server.stop() + end) + + it("should handle authentication state across server lifecycle", function() + -- Start with authentication + local success1, _ = server.start(config, test_auth_token) + expect(success1).to_be_true() + expect(server.state.auth_token).to_be(test_auth_token) + + -- Stop server + server.stop() + expect(server.state.auth_token).to_be_nil() + + -- Start without authentication + local success2, _ = server.start(config, nil) + expect(success2).to_be_true() + expect(server.state.auth_token).to_be_nil() + + -- Clean up + server.stop() + end) + + it("should handle different auth states", function() + -- Test with authentication enabled + local success1, _ = server.start(config, test_auth_token) + expect(success1).to_be_true() + expect(server.state.auth_token).to_be(test_auth_token) + + server.stop() + + -- Test with authentication disabled + local success2, _ = server.start(config, nil) + expect(success2).to_be_true() + expect(server.state.auth_token).to_be_nil() + + -- Clean up + server.stop() + end) + + it("should preserve auth token during handler setup", function() + -- Start server with auth token + server.start(config, test_auth_token) + expect(server.state.auth_token).to_be(test_auth_token) + + -- Register handlers - should not affect auth token + server.register_handlers() + expect(server.state.auth_token).to_be(test_auth_token) + + -- Get status - should not affect auth token + local status = server.get_status() + expect(status.running).to_be_true() + expect(server.state.auth_token).to_be(test_auth_token) + + -- Clean up + server.stop() + end) + + it("should handle multiple auth token operations", function() + -- Start server + server.start(config, test_auth_token) + expect(server.state.auth_token).to_be(test_auth_token) + + -- Multiple operations that should not affect auth token + for i = 1, 5 do + server.register_handlers() + local status = server.get_status() + expect(status.running).to_be_true() + + -- Auth token should remain stable + expect(server.state.auth_token).to_be(test_auth_token) + end + + -- Clean up + server.stop() + end) + end) + teardown() end) diff --git a/tests/lockfile_test.lua b/tests/lockfile_test.lua index 41a91b6..7ebbe51 100644 --- a/tests/lockfile_test.lua +++ b/tests/lockfile_test.lua @@ -24,12 +24,21 @@ if not _G.vim then fs = { remove = function() end }, ---@type vim_fs_module fn = { ---@type vim_fn_table expand = function(path) - return select(1, path:gsub("~", "/home/user")) + -- Use a temp directory that actually exists + local temp_dir = os.getenv("TMPDIR") or "/tmp" + return select(1, path:gsub("~", temp_dir .. "/claude_test")) end, -- Add other vim.fn mocks as needed by lockfile tests -- For now, only adding what's explicitly used or causing major type issues - filereadable = function() - return 1 + filereadable = function(path) + -- Check if file actually exists + local file = io.open(path, "r") + if file then + file:close() + return 1 + else + return 0 + end end, fnamemodify = function(fname, _) return fname @@ -113,8 +122,48 @@ if not _G.vim then }, }, json = { - encode = function(_obj) -- Prefix unused param with underscore - return '{"mocked":"json"}' + encode = function(obj) + -- Simple JSON encoding for testing + if type(obj) == "table" then + local pairs_array = {} + for k, v in pairs(obj) do + local key_str = '"' .. tostring(k) .. '"' + local val_str + if type(v) == "string" then + val_str = '"' .. v .. '"' + elseif type(v) == "number" then + val_str = tostring(v) + elseif type(v) == "table" then + -- Simple array encoding + local items = {} + for _, item in ipairs(v) do + table.insert(items, '"' .. tostring(item) .. '"') + end + val_str = "[" .. table.concat(items, ",") .. "]" + else + val_str = '"' .. tostring(v) .. '"' + end + table.insert(pairs_array, key_str .. ":" .. val_str) + end + return "{" .. table.concat(pairs_array, ",") .. "}" + else + return '"' .. tostring(obj) .. '"' + end + end, + decode = function(json_str) + -- Very basic JSON parsing for test purposes + if json_str:match("^%s*{.*}%s*$") then + local result = {} + -- Extract key-value pairs - this is very basic + for key, value in json_str:gmatch('"([^"]+)"%s*:%s*"([^"]*)"') do + result[key] = value + end + for key, value in json_str:gmatch('"([^"]+)"%s*:%s*(%d+)') do + result[key] = tonumber(value) + end + return result + end + return {} end, }, lsp = {}, -- Existing lsp mock part @@ -198,12 +247,22 @@ describe("Lockfile Module", function() return "/mock/cwd" end + -- Create test directory + local temp_dir = os.getenv("TMPDIR") or "/tmp" + local test_dir = temp_dir .. "/claude_test/.claude/ide" + os.execute("mkdir -p '" .. test_dir .. "'") + -- Load the lockfile module for all tests package.loaded["claudecode.lockfile"] = nil -- Clear any previous requires lockfile = require("claudecode.lockfile") end) teardown(function() + -- Clean up test files + local temp_dir = os.getenv("TMPDIR") or "/tmp" + local test_dir = temp_dir .. "/claude_test" + os.execute("rm -rf '" .. test_dir .. "'") + -- Restore original vim if real_vim then _G.vim = real_vim @@ -279,4 +338,63 @@ describe("Lockfile Module", function() assert(2 == #folders) -- cwd + 1 unique workspace folder end) end) + + describe("authentication token functionality", function() + it("should generate auth tokens", function() + local token1 = lockfile.generate_auth_token() + local token2 = lockfile.generate_auth_token() + + -- Tokens should be strings + assert("string" == type(token1)) + assert("string" == type(token2)) + + -- Tokens should be different + assert(token1 ~= token2) + + -- Tokens should match UUID format (xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx) + assert(token1:match("^%x+%-%x+%-4%x+%-[89ab]%x+%-%x+$")) + assert(token2:match("^%x+%-%x+%-4%x+%-[89ab]%x+%-%x+$")) + end) + + it("should create lock files with auth tokens", function() + local port = 12345 + local success, lock_path, auth_token = lockfile.create(port) + + assert(success == true) + assert("string" == type(lock_path)) + assert("string" == type(auth_token)) + + -- Should be able to read the auth token back + local read_success, read_token, read_error = lockfile.get_auth_token(port) + assert(read_success == true) + assert(auth_token == read_token) + assert(read_error == nil) + end) + + it("should create lock files with pre-generated auth tokens", function() + local port = 12346 + local preset_token = "test-auth-token-12345" + local success, lock_path, returned_token = lockfile.create(port, preset_token) + + assert(success == true) + assert("string" == type(lock_path)) + assert(preset_token == returned_token) + + -- Should be able to read the preset token back + local read_success, read_token, read_error = lockfile.get_auth_token(port) + assert(read_success == true) + assert(preset_token == read_token) + assert(read_error == nil) + end) + + it("should handle missing lock files when reading auth tokens", function() + local nonexistent_port = 99999 + local success, token, error = lockfile.get_auth_token(nonexistent_port) + + assert(success == false) + assert(token == nil) + assert("string" == type(error)) + assert(error:find("Lock file does not exist")) + end) + end) end) diff --git a/tests/unit/init_spec.lua b/tests/unit/init_spec.lua index 8840ea0..2149983 100644 --- a/tests/unit/init_spec.lua +++ b/tests/unit/init_spec.lua @@ -31,12 +31,15 @@ describe("claudecode.init", function() local mock_lockfile = { create = function() - return true, "/mock/path" + return true, "/mock/path", "mock-auth-token-12345" end, ---@type SpyableFunction remove = function() return true end, + generate_auth_token = function() + return "mock-auth-token-12345" + end, } local mock_selection = { diff --git a/tests/unit/server/handshake_spec.lua b/tests/unit/server/handshake_spec.lua new file mode 100644 index 0000000..23db897 --- /dev/null +++ b/tests/unit/server/handshake_spec.lua @@ -0,0 +1,279 @@ +require("tests.busted_setup") + +describe("WebSocket handshake authentication", function() + local handshake + + before_each(function() + handshake = require("claudecode.server.handshake") + end) + + after_each(function() + package.loaded["claudecode.server.handshake"] = nil + end) + + describe("validate_upgrade_request with authentication", function() + local valid_request_base = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "", + "", + }, "\r\n") + + it("should accept valid request with correct auth token", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + local request_with_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization: " .. expected_token, + "", + "", + }, "\r\n") + + local is_valid, headers = handshake.validate_upgrade_request(request_with_auth, expected_token) + + assert.is_true(is_valid) + assert.is_table(headers) + assert.equals(expected_token, headers["x-claude-code-ide-authorization"]) + end) + + it("should reject request with missing auth token when required", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + + local is_valid, error_msg = handshake.validate_upgrade_request(valid_request_base, expected_token) + + assert.is_false(is_valid) + assert.equals("Missing authentication header: x-claude-code-ide-authorization", error_msg) + end) + + it("should reject request with incorrect auth token", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + local wrong_token = "123e4567-e89b-12d3-a456-426614174000" + + local request_with_wrong_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization: " .. wrong_token, + "", + "", + }, "\r\n") + + local is_valid, error_msg = handshake.validate_upgrade_request(request_with_wrong_auth, expected_token) + + assert.is_false(is_valid) + assert.equals("Invalid authentication token", error_msg) + end) + + it("should accept request without auth token when none required", function() + local is_valid, headers = handshake.validate_upgrade_request(valid_request_base, nil) + + assert.is_true(is_valid) + assert.is_table(headers) + end) + + it("should reject request with empty auth token when required", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + + local request_with_empty_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization: ", + "", + "", + }, "\r\n") + + local is_valid, error_msg = handshake.validate_upgrade_request(request_with_empty_auth, expected_token) + + assert.is_false(is_valid) + assert.equals("Authentication token too short (min 10 characters)", error_msg) + end) + + it("should handle case-insensitive auth header name", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + + local request_with_uppercase_header = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "X-Claude-Code-IDE-Authorization: " .. expected_token, + "", + "", + }, "\r\n") + + local is_valid, headers = handshake.validate_upgrade_request(request_with_uppercase_header, expected_token) + + assert.is_true(is_valid) + assert.is_table(headers) + end) + end) + + describe("process_handshake with authentication", function() + local valid_request_base = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "", + "", + }, "\r\n") + + it("should complete handshake successfully with valid auth token", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + local request_with_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization: " .. expected_token, + "", + "", + }, "\r\n") + + local success, response, headers = handshake.process_handshake(request_with_auth, expected_token) + + assert.is_true(success) + assert.is_string(response) + assert.is_table(headers) + assert.matches("HTTP/1.1 101 Switching Protocols", response) + assert.matches("Upgrade: websocket", response) + assert.matches("Connection: Upgrade", response) + assert.matches("Sec%-WebSocket%-Accept:", response) + end) + + it("should fail handshake with missing auth token", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + + local success, response, headers = handshake.process_handshake(valid_request_base, expected_token) + + assert.is_false(success) + assert.is_string(response) + assert.is_nil(headers) + assert.matches("HTTP/1.1 400 Bad Request", response) + assert.matches("Missing authentication header", response) + end) + + it("should fail handshake with invalid auth token", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + local wrong_token = "123e4567-e89b-12d3-a456-426614174000" + + local request_with_wrong_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization: " .. wrong_token, + "", + "", + }, "\r\n") + + local success, response, headers = handshake.process_handshake(request_with_wrong_auth, expected_token) + + assert.is_false(success) + assert.is_string(response) + assert.is_nil(headers) + assert.matches("HTTP/1.1 400 Bad Request", response) + assert.matches("Invalid authentication token", response) + end) + + it("should complete handshake without auth when none required", function() + local success, response, headers = handshake.process_handshake(valid_request_base, nil) + + assert.is_true(success) + assert.is_string(response) + assert.is_table(headers) + assert.matches("HTTP/1.1 101 Switching Protocols", response) + end) + end) + + describe("authentication edge cases", function() + it("should handle malformed auth header format", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + + local request_with_malformed_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization:not-a-uuid", + "", + "", + }, "\r\n") + + local is_valid, error_msg = handshake.validate_upgrade_request(request_with_malformed_auth, expected_token) + + assert.is_false(is_valid) + assert.equals("Invalid authentication token", error_msg) + end) + + it("should handle multiple auth headers (uses last one)", function() + local expected_token = "550e8400-e29b-41d4-a716-446655440000" + local wrong_token = "123e4567-e89b-12d3-a456-426614174000" + + local request_with_multiple_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization: " .. wrong_token, + "x-claude-code-ide-authorization: " .. expected_token, + "", + "", + }, "\r\n") + + local is_valid, headers = handshake.validate_upgrade_request(request_with_multiple_auth, expected_token) + + assert.is_true(is_valid) + assert.is_table(headers) + assert.equals(expected_token, headers["x-claude-code-ide-authorization"]) + end) + + it("should reject very long auth tokens", function() + local expected_token = string.rep("a", 1000) -- Very long token + + local request_with_long_auth = table.concat({ + "GET /websocket HTTP/1.1", + "Host: localhost:8080", + "Upgrade: websocket", + "Connection: upgrade", + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", + "Sec-WebSocket-Version: 13", + "x-claude-code-ide-authorization: " .. expected_token, + "", + "", + }, "\r\n") + + local is_valid, error_msg = handshake.validate_upgrade_request(request_with_long_auth, expected_token) + + assert.is_false(is_valid) + assert.equals("Authentication token too long (max 500 characters)", error_msg) + end) + end) +end) diff --git a/tests/unit/server_spec.lua b/tests/unit/server_spec.lua index 5bfc81c..d2df77a 100644 --- a/tests/unit/server_spec.lua +++ b/tests/unit/server_spec.lua @@ -170,6 +170,142 @@ describe("WebSocket Server", function() server.stop() end) + describe("authentication integration", function() + it("should start server with authentication token", function() + local config = { + port_range = { + min = 10000, + max = 65535, + }, + } + local auth_token = "550e8400-e29b-41d4-a716-446655440000" + + local success, port = server.start(config, auth_token) + + expect(success).to_be_true() + expect(server.state.auth_token).to_be(auth_token) + expect(server.state.server).to_be_table() + expect(server.state.port).to_be(port) + + -- Clean up + server.stop() + end) + + it("should clear auth token when server stops", function() + local config = { + port_range = { + min = 10000, + max = 65535, + }, + } + local auth_token = "550e8400-e29b-41d4-a716-446655440000" + + -- Start server with auth token + server.start(config, auth_token) + expect(server.state.auth_token).to_be(auth_token) + + -- Stop server + server.stop() + expect(server.state.auth_token).to_be_nil() + end) + + it("should start server without authentication token", function() + local config = { + port_range = { + min = 10000, + max = 65535, + }, + } + + local success, port = server.start(config, nil) + + expect(success).to_be_true() + expect(server.state.auth_token).to_be_nil() + expect(server.state.server).to_be_table() + expect(server.state.port).to_be(port) + + -- Clean up + server.stop() + end) + + it("should pass auth token to TCP server creation", function() + local config = { + port_range = { + min = 10000, + max = 65535, + }, + } + local auth_token = "550e8400-e29b-41d4-a716-446655440000" + + -- Mock the TCP server module to verify auth token is passed + local tcp_server = require("claudecode.server.tcp") + local original_create_server = tcp_server.create_server + local captured_auth_token = nil + + tcp_server.create_server = function(cfg, callbacks, auth_token_arg) + captured_auth_token = auth_token_arg + return original_create_server(cfg, callbacks, auth_token_arg) + end + + local success, _ = server.start(config, auth_token) + + -- Restore original function + tcp_server.create_server = original_create_server + + expect(success).to_be_true() + expect(captured_auth_token).to_be(auth_token) + + -- Clean up + server.stop() + end) + + it("should maintain auth token in server state throughout lifecycle", function() + local config = { + port_range = { + min = 10000, + max = 65535, + }, + } + local auth_token = "550e8400-e29b-41d4-a716-446655440000" + + -- Start server + server.start(config, auth_token) + expect(server.state.auth_token).to_be(auth_token) + + -- Get status should show running state + local status = server.get_status() + expect(status.running).to_be_true() + expect(server.state.auth_token).to_be(auth_token) + + -- Send message should work with auth token in place + local client = { id = "test_client" } + local success = server.send(client, "test_method", { test = "data" }) + expect(success).to_be_true() + expect(server.state.auth_token).to_be(auth_token) + + -- Clean up + server.stop() + end) + + it("should reject starting server if auth token is explicitly false", function() + local config = { + port_range = { + min = 10000, + max = 65535, + }, + } + + -- Use an empty string as invalid auth token + local success, _ = server.start(config, "") + + expect(success).to_be_true() -- Server should still start, just with empty token + expect(server.state.auth_token).to_be("") + + -- Clean up + server.stop() + end) + end) + -- Clean up after all tests teardown() end) From 91357d810ccf92f6169f3754436901c6ff5237ec Mon Sep 17 00:00:00 2001 From: Kohei Watanabe <31376689+nabekou29@users.noreply.github.com> Date: Sat, 21 Jun 2025 23:46:25 +0900 Subject: [PATCH 4/4] feat: add CLAUDE_CONFIG_DIR environment variable support (#58) --- README.md | 4 ++-- lua/claudecode/lockfile.lua | 12 +++++++++++- scripts/lib_claude.sh | 8 ++++++-- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 358ebe9..b3c31c7 100644 --- a/README.md +++ b/README.md @@ -112,7 +112,7 @@ This plugin creates a WebSocket server that Claude Code CLI connects to, impleme The protocol uses a WebSocket-based variant of MCP (Model Context Protocol) that: 1. Creates a WebSocket server on a random port -2. Writes a lock file to `~/.claude/ide/[port].lock` with connection info +2. Writes a lock file to `~/.claude/ide/[port].lock` (or `$CLAUDE_CONFIG_DIR/ide/[port].lock` if `CLAUDE_CONFIG_DIR` is set) with connection info 3. Sets environment variables that tell Claude where to connect 4. Implements MCP tools that Claude can call @@ -176,7 +176,7 @@ For deep technical details, see [ARCHITECTURE.md](./ARCHITECTURE.md). ## Troubleshooting -- **Claude not connecting?** Check `:ClaudeCodeStatus` and verify lock file exists in `~/.claude/ide/` +- **Claude not connecting?** Check `:ClaudeCodeStatus` and verify lock file exists in `~/.claude/ide/` (or `$CLAUDE_CONFIG_DIR/ide/` if `CLAUDE_CONFIG_DIR` is set) - **Need debug logs?** Set `log_level = "debug"` in opts - **Terminal issues?** Try `provider = "native"` if using snacks.nvim diff --git a/lua/claudecode/lockfile.lua b/lua/claudecode/lockfile.lua index 1590dad..6c19483 100644 --- a/lua/claudecode/lockfile.lua +++ b/lua/claudecode/lockfile.lua @@ -7,7 +7,17 @@ local M = {} --- Path to the lock file directory -M.lock_dir = vim.fn.expand("~/.claude/ide") +---@return string lock_dir The path to the lock file directory +local function get_lock_dir() + local claude_config_dir = os.getenv("CLAUDE_CONFIG_DIR") + if claude_config_dir and claude_config_dir ~= "" then + return vim.fn.expand(claude_config_dir .. "/ide") + else + return vim.fn.expand("~/.claude/ide") + end +end + +M.lock_dir = get_lock_dir() -- Track if random seed has been initialized local random_initialized = false diff --git a/scripts/lib_claude.sh b/scripts/lib_claude.sh index 486474a..c378203 100755 --- a/scripts/lib_claude.sh +++ b/scripts/lib_claude.sh @@ -3,7 +3,11 @@ # This library provides reusable functions for interacting with Claude Code's WebSocket API # Configuration -export CLAUDE_LOCKFILE_DIR="$HOME/.claude/ide" +if [ -n "$CLAUDE_CONFIG_DIR" ]; then + export CLAUDE_LOCKFILE_DIR="$CLAUDE_CONFIG_DIR/ide" +else + export CLAUDE_LOCKFILE_DIR="$HOME/.claude/ide" +fi export CLAUDE_LOG_DIR="mcp_test_logs" # Default log directory export CLAUDE_WS_TIMEOUT=10 # Default timeout in seconds @@ -13,7 +17,7 @@ export CLAUDE_WS_TIMEOUT=10 # Default timeout in seconds # Find the Claude lockfile and extract the port find_claude_lockfile() { # Get all .lock files - lock_files=$(find ~/.claude/ide -name "*.lock" 2>/dev/null || echo "") + lock_files=$(find "$CLAUDE_LOCKFILE_DIR" -name "*.lock" 2>/dev/null || echo "") if [ -z "$lock_files" ]; then echo "No Claude lockfiles found. Is the VSCode extension running?" >&2