diff --git a/.env.example b/.env.example index cd0d11c..2a65441 100644 --- a/.env.example +++ b/.env.example @@ -72,6 +72,10 @@ OPENROUTER_API_KEY= # Web search BRAVE_API_KEY= +# Paperclip — biomedical paper search (8M+ papers from bioRxiv, medRxiv, PMC, arXiv) +# Get API key at https://paperclip.gxl.ai +# PAPERCLIP_API_KEY= + # GitHub integration GITHUB_TOKEN= diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d4c23e..19241f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,62 @@ # Changelog +## v0.6.0 + +MCP server fixes, per-server mode configuration, @ mention system for referencing resources, parallel file inspection tool, conversation title generation improvements, and security hardening. + +### MCP Server Fixes +- **Celery worker MCP loading** -- MCP servers are now loaded in the Celery background worker path, matching the inline session path. Previously MCP tools were completely unavailable when using background jobs +- **Multi-server dispatch** -- Tools from multiple MCP servers now dispatch to their correct originating client. Previously `self._mcp_client` was overwritten by each server, causing tools from earlier servers to fail at execution time +- **Plan mode MCP access** -- MCP tools are no longer blocked by the plan-mode whitelist. Each tool is tracked with its own mode configuration and bypasses the built-in tool restrictions +- **Exception logging** -- `register_mcp_tools` now logs warnings on failure instead of silently swallowing exceptions with bare `except: pass` +- **Tool name collision logging** -- MCP tools that attempt to shadow built-in tool names are logged with a warning for security observability +- **Connection timeout** -- MCP server connections in the Celery worker are wrapped with a 30-second timeout to prevent hanging workers +- **Cleanup** -- MCP connections are properly disconnected in the Celery worker's finally block + +### Per-Server Mode Configuration +- **Mode checkboxes** -- Each MCP server can be configured to be available in Plan mode, Execute mode, or both (default: both) via checkboxes in Settings > MCP Servers +- **Backend enforcement** -- The `modes` field is stored in the server config, passed through `MCPManager.connect_servers` to `register_mcp_tools`, and enforced by `ToolRouter.is_tool_allowed` per-tool +- **Status endpoint** -- `GET /api/mcp/status` now includes `modes` in each server's response + +### @ Mention System +- **MentionPopover component** -- Type `@` in the chat input to open a dropdown showing MCP servers and workspace files. Supports directory browsing (typing `@code/` lists files in `code/`), keyboard navigation (arrows, Enter, Tab, Escape), and filtering by name +- **Mention chips** -- Active mentions are displayed as colored chips above the input area (blue for MCP servers, amber for files) +- **Resource references** -- Mentions are sent as lightweight structured references (`{type, value}`) alongside the message. The backend prepends reference hints that instruct the agent to use appropriate tools (`read`, `inspect_files`, MCP tools) to interact with the referenced resources +- **Mention model** -- New `Mention` Pydantic model with `type: Literal["server", "file"]` and `value: str` (max 1024 chars). Added `mentions` field to `MessageSend` +- **Input sanitization** -- Mention values are sanitized (backticks, newlines, control characters stripped, length capped at 256) before interpolation into prompt text to prevent LLM prompt injection + +### inspect_files Tool +- **Parallel file reading** -- New `inspect_files` tool reads multiple files or directories concurrently via `asyncio.gather` and scores each file for relevance against a user query +- **Keyword relevance scoring** -- Files are scored by keyword overlap between their content and the query, sorted by relevance, and returned within a configurable token budget (100K chars default) +- **Directory expansion** -- Directory paths are expanded to their file listings; hidden files (dotfiles) are excluded +- **Safety limits** -- Max 50 files per call, 200 lines per file for scoring, 2MB file-size gate (large files skipped before reading), negative `max_files` clamped to 1 +- **Security** -- Each child file in expanded directories is re-validated via `_validate_path` to catch symlinks escaping the workspace +- **Plan mode access** -- Added to the plan-mode allowlist for read-only context gathering + +### Conversation Title Generation +- **Deferred generation** -- Title generation no longer triggers after the 1st user message. It now triggers after the 3rd user message or on page refresh, whichever comes first +- **No re-updates** -- Once a title is set, it is not overwritten by subsequent triggers. A race-condition guard in `_auto_title` re-checks the current title from DB before persisting +- **Trigger guard** -- The `send_message` endpoint checks `conv.title == "New conversation"` before triggering, preventing redundant generation + +### Security Hardening +- Symlink traversal protection in `inspect_files` -- each child entry in expanded directories is validated via `_validate_path` before reading +- File-size gate in `inspect_files` -- files over 2MB are skipped before `read_text()` to prevent OOM +- `asyncio.get_running_loop()` used instead of deprecated `get_event_loop()` in async contexts +- MCP tool name shadowing logged as a warning for security observability +- Mention values sanitized to strip backticks, newlines, and control characters before prompt interpolation +- `Mention.value` field constrained to max 1024 characters via Pydantic `Field` +- MCP connection timeout (30s) in Celery worker prevents indefinite worker stalls + +### UI Fixes +- **Layout gap fix** -- Fixed 1px gap between the main content area and the right sidebar caused by `paddingRight` being 1px larger than the RightPanel's rendered width (`289px` → `288px`, `49px` → `48px`) +- **MCP live connection status** -- MCP server dots in the right panel now turn green when connected. Previously the status was hardcoded to `connected: false` in the REST endpoint. An `mcp_status` SSE event is now broadcast from both the session manager and Celery worker after `MCPManager.connect_servers()` succeeds, and the frontend handles it to update the dots in real time +- **Pre-existing lint fixes** -- Removed extraneous f-string prefix in `papers.py`, removed unused `AsyncMock` import in `test_tools_papers.py` + +### Testing +- **34 new backend tests** -- MCP multi-client dispatch (9), inspect_files tool (12), mention enrichment (7), Mention model validation (3), title generation (3 from prior session) +- **Total: 915 backend + 223 frontend = 1,138 tests** +- All ruff checks pass, frontend eslint 0 errors + ## v0.5.0 Project-scoped conversations, unified file workspace, Monaco code viewer, TODO approval flow, comprehensive agent guidance, and test infrastructure improvements. diff --git a/README.md b/README.md index e6817e6..bac8241 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,8 @@ - **Background jobs** — Celery + Redis. Close the browser, come back later. - **Multi-provider LLMs** — OpenAI, Anthropic, OpenRouter, plus local models (Ollama, LM Studio). Add custom providers with OpenAI SDK, Anthropic SDK, OpenRouter, or LiteLLM compatibility. - **Model picker** — Browse models grouped by provider with logos, sorted by release date. Recently used models at the top. Fetches live from [models.dev](https://models.dev). -- **MCP servers** — Connect remote HTTP/HTTPS MCP servers with custom authentication (Bearer, API key, headers). +- **MCP servers** — Connect remote HTTP/HTTPS MCP servers with custom authentication (Bearer, API key, headers). Configure per-server mode availability (Plan, Execute, or both). Live connection status in the sidebar. +- **@ mentions** — Type `@` in the chat to reference MCP servers or workspace files/directories. The agent uses its tools to interact with the referenced resources. - **Onboarding flow** — Guided setup when no LLM provider is configured. ## Quick Start diff --git a/backend/openmlr/models.py b/backend/openmlr/models.py index 35c8870..3c09287 100644 --- a/backend/openmlr/models.py +++ b/backend/openmlr/models.py @@ -70,12 +70,23 @@ class ConversationDetail(BaseModel): # ---- Messaging ---- +class Mention(BaseModel): + """A resource reference from an @ mention in the chat input.""" + + type: Literal["server", "file"] + value: str = Field( + max_length=1024, + description="server name or workspace-relative file/directory path", + ) + + class MessageSend(BaseModel): message: str mode: Literal["plan", "execute"] | None = ( None # per-message mode; only plan or execute accepted ) request_id: str | None = None # client-generated idempotency key + mentions: list[Mention] | None = None # @ mention references class ApprovalRequest(BaseModel): diff --git a/backend/openmlr/routes/agent.py b/backend/openmlr/routes/agent.py index e364d02..5d440e7 100644 --- a/backend/openmlr/routes/agent.py +++ b/backend/openmlr/routes/agent.py @@ -2,6 +2,7 @@ import asyncio import logging +import re from typing import Annotated from fastapi import APIRouter, Depends, HTTPException, Request @@ -19,6 +20,12 @@ router = APIRouter(prefix="/api", tags=["agent"]) +# Default conversation title — used as sentinel to detect untitled conversations +_DEFAULT_TITLE = "New conversation" + +# Regex for stripping control characters from mention values (prompt injection defense) +_MENTION_SANITIZE_RE = re.compile(r"[`\x00-\x1f]") + def _sm(request: Request): return request.app.state.session_manager @@ -147,8 +154,8 @@ async def get_conversation( conv = await _get_conv_or_404(db, uuid, user.id) msgs = await ops.get_messages(db, conv.id) - # Re-generate title if still "New conversation" and has messages - if conv.title == "New conversation" and msgs: + # Re-generate title if still the default and has messages + if conv.title == _DEFAULT_TITLE and msgs: msg_dicts = [_msg_dict(m) for m in msgs] _task = asyncio.create_task( _auto_title(_sm(request), _bus(request), db, conv.id, conv.uuid, msg_dicts) @@ -420,7 +427,10 @@ async def send_message( # If conversation has no model, use user's default effective_model = conv.model or user_default_model - # Title generation after 1st and 3rd messages + # Enrich message with @ mention reference hints + enriched_message = _enrich_with_mentions(body.message, body.mentions) + + # Title generation after 3rd user message (if not already titled) user_count = (conv.user_message_count or 0) + 1 # If background jobs are enabled, use Celery @@ -429,14 +439,14 @@ async def send_message( db=db, conversation_id=conv.id, user_id=user.id, - message=body.message, + message=enriched_message, mode=body.mode, model=effective_model, uuid=conv.uuid, ) # Title generation (still async in web process for now) - if user_count in (1, 3): + if user_count == 3 and conv.title == _DEFAULT_TITLE: msg_dicts = await _load_messages(db, conv.id) _task = asyncio.create_task( _auto_title(sm, event_bus, db, conv.id, conv.uuid, msg_dicts) @@ -445,7 +455,7 @@ async def send_message( return {"ok": True, "job_id": job.job_id if job else None, "background": True} # Synchronous processing (original flow) - # Persist user message to DB + # Persist original message to DB (without enrichment clutter) await ops.add_message(db, conv.id, "user", body.message) await ops.increment_user_message_count(db, conv.id) @@ -475,9 +485,9 @@ async def send_message( _wire_persistence(active, db, conv.id) active._persist_wired = True - _task = asyncio.create_task(sm.process_message(conv.id, body.message, mode=body.mode)) + _task = asyncio.create_task(sm.process_message(conv.id, enriched_message, mode=body.mode)) - if user_count in (1, 3): + if user_count == 3 and conv.title == _DEFAULT_TITLE: msg_dicts = await _load_messages(db, conv.id) _task = asyncio.create_task(_auto_title(sm, event_bus, db, conv.id, conv.uuid, msg_dicts)) @@ -753,6 +763,46 @@ def _conv_dict(c) -> dict: } +def _sanitize_mention_value(v: str) -> str: + """Strip control characters and cap length to prevent prompt injection.""" + v = v[:256] + return _MENTION_SANITIZE_RE.sub("", v) + + +def _enrich_with_mentions(message: str, mentions: list | None) -> str: + """Prepend resource-reference hints for @ mentions. + + Mentions are lightweight pointers — the agent is expected to use its + tools (``read``, ``inspect_files``, MCP tools) to interact with them. + """ + if not mentions: + return message + + refs: list[str] = [] + for m in mentions: + safe_value = _sanitize_mention_value(m.value) + if m.type == "file": + path = safe_value.rstrip("/") + if m.value.endswith("/"): + refs.append( + f"- Directory {path}/ — list its contents with read or use " + f"inspect_files to inspect relevant files." + ) + else: + refs.append(f"- File {path} — use read to inspect this file.") + elif m.type == "server": + refs.append(f"- MCP Server {safe_value} — use tools provided by this server.") + + if not refs: + return message + + hint = ( + "[The user referenced these resources — use the appropriate tools " + "to interact with them before responding:]\n" + "\n".join(refs) + ) + return hint + "\n\n" + message + + def _msg_dict(m) -> dict: return { "id": m.id, @@ -801,6 +851,13 @@ async def _auto_title(sm, event_bus, db, conv_id, uuid, messages): title = await LLMProvider.generate_title(messages, config) if title: + # Re-check the current title to avoid overwriting a title + # that was already set by another trigger (e.g. page refresh). + current_conv = await ops.get_conversation_by_id(db, conv_id) + if current_conv and current_conv.title != _DEFAULT_TITLE: + logger.debug(f"Skipping title update for conv {conv_id}: already titled") + return + await ops.update_conversation_title(db, conv_id, title) await event_bus.broadcast( AgentEvent(event_type="conversation_updated", data={"uuid": uuid, "title": title}) diff --git a/backend/openmlr/routes/mcp.py b/backend/openmlr/routes/mcp.py index 7237dab..9cff9c4 100644 --- a/backend/openmlr/routes/mcp.py +++ b/backend/openmlr/routes/mcp.py @@ -58,6 +58,7 @@ async def get_status( "url": config.get("url", ""), "enabled": config.get("enabled", True), "connected": False, # Will be updated via SSE in real-time + "modes": config.get("modes", ["plan", "execute"]), } ) diff --git a/backend/openmlr/routes/settings.py b/backend/openmlr/routes/settings.py index ce30d3d..15625cb 100644 --- a/backend/openmlr/routes/settings.py +++ b/backend/openmlr/routes/settings.py @@ -62,6 +62,7 @@ async def update_setting( "github_token": "GITHUB_TOKEN", "semantic_scholar_api_key": "SEMANTIC_SCHOLAR_API_KEY", "openalex_api_key": "OPENALEX_API_KEY", + "paperclip_api_key": "PAPERCLIP_API_KEY", "modal_token_id": "MODAL_TOKEN_ID", "modal_token_secret": "MODAL_TOKEN_SECRET", "hf_token": "HF_TOKEN", @@ -127,6 +128,7 @@ def _is_provider_configured(provider_id: str, provider_settings: dict) -> bool: "github": "GITHUB_TOKEN", "semantic_scholar": "SEMANTIC_SCHOLAR_API_KEY", "openalex": "OPENALEX_API_KEY", + "paperclip": "PAPERCLIP_API_KEY", "modal": "MODAL_TOKEN_ID", "huggingface": "HF_TOKEN", } @@ -144,6 +146,7 @@ def _is_provider_configured(provider_id: str, provider_settings: dict) -> bool: "github": "github_token", "semantic_scholar": "semantic_scholar_api_key", "openalex": "openalex_api_key", + "paperclip": "paperclip_api_key", "modal": "modal_token_id", "huggingface": "hf_token", }.get(provider_id) @@ -261,6 +264,14 @@ async def list_providers( "categories": ["compute"], "docs_url": "https://modal.com/docs", }, + { + "id": "paperclip", + "name": "Paperclip", + "key_env": "PAPERCLIP_API_KEY", + "configured": _is_provider_configured("paperclip", provider_settings), + "categories": ["papers"], + "docs_url": "https://paperclip.gxl.ai/docs", + }, { "id": "huggingface", "name": "Hugging Face", @@ -768,6 +779,7 @@ async def save_config( "GITHUB_TOKEN", "SEMANTIC_SCHOLAR_API_KEY", "OPENALEX_API_KEY", + "PAPERCLIP_API_KEY", "MODAL_TOKEN_ID", "MODAL_TOKEN_SECRET", "HF_TOKEN", diff --git a/backend/openmlr/services/session_manager.py b/backend/openmlr/services/session_manager.py index ccdc255..6d04692 100644 --- a/backend/openmlr/services/session_manager.py +++ b/backend/openmlr/services/session_manager.py @@ -191,6 +191,13 @@ async def get_or_create_session( ) if count > 0: log.info(f"Session {conversation_id}: loaded {count} MCP tools") + # Broadcast live connection status to the frontend + await self.event_bus.broadcast( + AgentEvent( + event_type="mcp_status", + data={"servers": mcp_manager.get_server_statuses()}, + ) + ) except Exception as e: log.warning(f"Session {conversation_id}: failed to load MCP servers - {e}") diff --git a/backend/openmlr/tasks/agent_tasks.py b/backend/openmlr/tasks/agent_tasks.py index 277925f..3b3fc3d 100644 --- a/backend/openmlr/tasks/agent_tasks.py +++ b/backend/openmlr/tasks/agent_tasks.py @@ -128,6 +128,9 @@ async def _async_process_message( sandbox_manager = SandboxManager() tool_router = create_tool_router(sandbox_manager) + # Track MCP manager for cleanup in finally block + mcp_manager = None + # Resolve project workspace for workspace tools and local tools async with worker_session() as db: try: @@ -148,6 +151,33 @@ async def _async_process_message( except Exception as e: logger.warning(f"Worker job {job_id}: failed to resolve project workspace - {e}") + # Load MCP servers from user settings (with timeout to avoid stalling the worker) + async with worker_session() as db: + try: + from ..tools.mcp import MCPManager + + user_settings = await ops.get_all_settings(db, user_id, category="mcp") + mcp_servers = user_settings.get("mcp", {}).get("servers", {}) + if mcp_servers: + mcp_manager = MCPManager() + count = await asyncio.wait_for( + mcp_manager.connect_servers(mcp_servers, tool_router, blocklist=set()), + timeout=30.0, + ) + if count > 0: + logger.info(f"Worker job {job_id}: loaded {count} MCP tools") + # Broadcast live connection status to frontend + await publish_event( + AgentEvent( + event_type="mcp_status", + data={"servers": mcp_manager.get_server_statuses()}, + ) + ) + except TimeoutError: + logger.warning(f"Worker job {job_id}: MCP server connection timed out") + except Exception as e: + logger.warning(f"Worker job {job_id}: failed to load MCP servers - {e}") + # Build and set system prompt session.context_manager.system_prompt = build_system_prompt( tool_specs=tool_router.get_raw_specs(), @@ -264,6 +294,13 @@ async def _poll_interrupt(): except Exception: pass + # Disconnect MCP servers + if mcp_manager: + try: + await mcp_manager.disconnect_all() + except Exception: + pass + # Clear any lingering interrupt key try: from ..services.redis_pubsub import clear_interrupt diff --git a/backend/openmlr/tools/inspect.py b/backend/openmlr/tools/inspect.py new file mode 100644 index 0000000..2bfca7d --- /dev/null +++ b/backend/openmlr/tools/inspect.py @@ -0,0 +1,274 @@ +"""inspect_files tool — parallel file reading with relevance filtering.""" + +import asyncio +import logging +from pathlib import Path + +from ..agent.types import ToolSpec + +logger = logging.getLogger("openmlr.tools.inspect") + +# Limits +_MAX_FILES = 50 # max files to read in one call +_MAX_LINES_PER_FILE = 200 # lines per file for relevance check +_MAX_TOTAL_CHARS = 100_000 # total output budget +_MAX_FILE_SIZE = 2 * 1024 * 1024 # 2 MB — skip files larger than this + + +def create_inspect_tool() -> ToolSpec: + return ToolSpec( + name="inspect_files", + description=( + "Read multiple files or directories in parallel and filter for relevance.\n\n" + "Use this when you need to quickly scan many files (e.g. an entire directory) " + "to find which ones are relevant to the user's request. Each file is read " + "concurrently and scored for relevance against the query.\n\n" + "Returns relevant file contents and a list of skipped files with reasons.\n\n" + "Examples:\n" + '- inspect_files(paths=["code/"], query="training loop")\n' + '- inspect_files(paths=["code/train.py", "code/model.py"], query="loss function")\n' + '- inspect_files(paths=["data/"], query="dataset loading")' + ), + parameters={ + "type": "object", + "properties": { + "paths": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "File or directory paths to inspect. Directories are expanded " + "to their file listing. Relative paths resolve from workspace root." + ), + }, + "query": { + "type": "string", + "description": ( + "What you're looking for — used to filter which files are relevant. " + "Be specific (e.g. 'training loop implementation' not just 'training')." + ), + }, + "max_files": { + "type": "integer", + "description": f"Max files to read (default {_MAX_FILES}).", + }, + }, + "required": ["paths", "query"], + }, + handler=_handle_inspect_files, + ) + + +# --------------------------------------------------------------------------- +# Phase helpers (extracted to reduce cognitive complexity of the main handler) +# --------------------------------------------------------------------------- + + +def _expand_paths( + paths: list[str], effective_root: Path, validate_path_fn +) -> tuple[list[Path], list[str]]: + """Expand user-supplied paths into validated file paths and directory listings.""" + file_paths: list[Path] = [] + dir_listings: list[str] = [] + + for p in paths: + target = Path(p).expanduser() + if not target.is_absolute(): + target = effective_root / target + target, error = validate_path_fn(target) + if error: + dir_listings.append(f"- `{p}`: {error}") + continue + + if target.is_dir(): + _expand_directory(target, effective_root, validate_path_fn, file_paths, dir_listings) + elif target.is_file(): + file_paths.append(target) + else: + dir_listings.append(f"- `{p}`: not found") + + return file_paths, dir_listings + + +def _expand_directory( + target: Path, + effective_root: Path, + validate_path_fn, + file_paths: list[Path], + dir_listings: list[str], +) -> None: + """Expand a single directory into file_paths and dir_listings (mutates in place).""" + entries = sorted(target.iterdir()) + dir_entries: list[str] = [] + for entry in entries: + if entry.name.startswith("."): + continue + if entry.is_file(): + resolved_entry, entry_err = validate_path_fn(entry.resolve()) + if not entry_err: + file_paths.append(resolved_entry) + dir_entries.append(f" {'/' if entry.is_dir() else ' '} {entry.name}") + rel = target.relative_to(effective_root) if _is_relative(target, effective_root) else target + dir_listings.append(f"Directory `{rel}/`:\n" + "\n".join(dir_entries)) + + +def _score_and_filter( + file_paths: list[Path], + read_results: list, + effective_root: Path, + query: str, +) -> tuple[list[dict], list[dict]]: + """Score files for relevance and split into relevant/skipped lists.""" + file_data: list[dict] = [] + for fp, result in zip(file_paths, read_results, strict=True): + if isinstance(result, Exception): + continue + content, line_count = result + if not content.strip(): + continue + rel_path = fp.relative_to(effective_root) if _is_relative(fp, effective_root) else fp + score = _score_relevance(content, query) + file_data.append( + {"path": str(rel_path), "content": content, "lines": line_count, "score": score} + ) + + file_data.sort(key=lambda x: x["score"], reverse=True) + + relevant: list[dict] = [] + skipped: list[dict] = [] + total_chars = 0 + for fd in file_data: + if fd["score"] < 0.1 or total_chars + len(fd["content"]) > _MAX_TOTAL_CHARS: + skipped.append(fd) + else: + relevant.append(fd) + total_chars += len(fd["content"]) + + return relevant, skipped + + +def _format_output( + dir_listings: list[str], + relevant: list[dict], + skipped: list[dict], + total_inspected: int, +) -> str: + """Build the final markdown output string.""" + parts: list[str] = [] + + if dir_listings: + parts.append("## Directory Listings\n" + "\n".join(dir_listings)) + + parts.append(f"## Relevant Files ({len(relevant)}/{total_inspected} inspected)") + for fd in relevant: + parts.append( + f"\n### {fd['path']} ({fd['lines']} lines, relevance: {fd['score']:.0%})\n" + f"```\n{fd['content']}\n```" + ) + + if skipped: + skip_lines = [f"- `{fd['path']}` ({fd['lines']} lines)" for fd in skipped] + parts.append( + f"\n## Skipped ({len(skipped)} files — low relevance or budget exceeded)\n" + + "\n".join(skip_lines) + ) + + return "\n\n".join(parts) + + +# --------------------------------------------------------------------------- +# Main handler +# --------------------------------------------------------------------------- + + +async def _handle_inspect_files( + paths: list[str], + query: str, + max_files: int = _MAX_FILES, + **kwargs, +) -> tuple[str, bool]: + """Read files in parallel, score for relevance, return filtered results.""" + from .local import _get_effective_root, _validate_path + + if not paths: + return "No paths provided.", False + if not query: + return "No query provided — specify what you're looking for.", False + + effective_root = _get_effective_root() + max_files = max(1, min(max_files, _MAX_FILES)) + + # Phase 1: Expand directories to file lists + file_paths, dir_listings = _expand_paths(paths, effective_root, _validate_path) + + if len(file_paths) > max_files: + skipped_count = len(file_paths) - max_files + file_paths = file_paths[:max_files] + dir_listings.append( + f"\n[Truncated: only inspecting first {max_files} files, " + f"{skipped_count} additional files skipped]" + ) + + if not file_paths: + result = "No files found to inspect.\n" + if dir_listings: + result += "\n".join(dir_listings) + return result, True + + # Phase 2: Read all files in parallel + loop = asyncio.get_running_loop() + read_tasks = [loop.run_in_executor(None, _read_file_snippet, fp) for fp in file_paths] + read_results = await asyncio.gather(*read_tasks, return_exceptions=True) + + # Phase 3: Score relevance and filter + relevant, skipped = _score_and_filter(file_paths, read_results, effective_root, query) + + # Phase 4: Format output + return _format_output(dir_listings, relevant, skipped, len(relevant) + len(skipped)), True + + +# --------------------------------------------------------------------------- +# Utility functions +# --------------------------------------------------------------------------- + + +def _read_file_snippet(path: Path) -> tuple[str, int]: + """Read up to _MAX_LINES_PER_FILE lines from a file. Returns (content, total_lines).""" + try: + size = path.stat().st_size + if size > _MAX_FILE_SIZE: + return f"[File too large: {size:,} bytes, skipped]", 0 + text = path.read_text(encoding="utf-8", errors="replace") + except Exception: + return "", 0 + + lines = text.splitlines() + total = len(lines) + selected = lines[:_MAX_LINES_PER_FILE] + content = "\n".join(f"{i}: {line}" for i, line in enumerate(selected, 1)) + if total > _MAX_LINES_PER_FILE: + content += f"\n\n[... {total - _MAX_LINES_PER_FILE} more lines truncated]" + return content, total + + +def _score_relevance(content: str, query: str) -> float: + """Simple keyword-overlap relevance score (0.0 to 1.0). + + Splits the query into terms and checks how many appear in the content. + This is fast and doesn't require an LLM call. + """ + content_lower = content.lower() + terms = [t.strip() for t in query.lower().split() if len(t.strip()) > 2] + if not terms: + return 0.5 # no useful query terms — include by default + + matches = sum(1 for t in terms if t in content_lower) + return matches / len(terms) + + +def _is_relative(path: Path, root: Path) -> bool: + """Check if path is under root without raising.""" + try: + path.relative_to(root) + return True + except ValueError: + return False diff --git a/backend/openmlr/tools/mcp.py b/backend/openmlr/tools/mcp.py index 87bd888..8bdf4f5 100644 --- a/backend/openmlr/tools/mcp.py +++ b/backend/openmlr/tools/mcp.py @@ -140,9 +140,11 @@ async def connect_servers( # Connect and register tools await client.__aenter__() + modes = config.get("modes", ["plan", "execute"]) count = await tool_router.register_mcp_tools( client, blocklist=blocklist or set(), + modes=modes, ) self._clients[server_name] = client diff --git a/backend/openmlr/tools/papers.py b/backend/openmlr/tools/papers.py index 0079d18..c54fda9 100644 --- a/backend/openmlr/tools/papers.py +++ b/backend/openmlr/tools/papers.py @@ -23,6 +23,7 @@ ARXIV_API = "https://export.arxiv.org/api/query" AR5IV_BASE = "https://ar5iv.labs.arxiv.org/html" PWC_API = "https://paperswithcode.com/api/v1" +PAPERCLIP_API = "https://paperclip.gxl.ai" # OpenAlex: API key or polite pool via mailto MAILTO = os.environ.get("OPENALEX_EMAIL", "openmlr@example.com") @@ -60,9 +61,11 @@ def create_papers_tool() -> ToolSpec: return ToolSpec( name="papers", description=( - "Search and read academic papers using OpenAlex, Semantic Scholar, arXiv, CrossRef, and Papers With Code. " + "Search and read academic papers using OpenAlex, Semantic Scholar, arXiv, CrossRef, " + "Papers With Code, and Paperclip (8M+ biomedical papers from bioRxiv, medRxiv, PMC). " "Multi-source search with automatic fallback for best results. " "Operations: search (OpenAlex+S2), arxiv_search (arXiv direct), semantic_search (Semantic Scholar), " + "paperclip_search (biomedical: bioRxiv/medRxiv/PMC/arXiv), paperclip_lookup (lookup by DOI/PMID), " "trending, details, read_paper, citations, recommend, find_code, find_datasets, " "author_papers." ), @@ -75,6 +78,8 @@ def create_papers_tool() -> ToolSpec: "search", "arxiv_search", "semantic_search", + "paperclip_search", + "paperclip_lookup", "trending", "details", "read_paper", @@ -89,6 +94,8 @@ def create_papers_tool() -> ToolSpec: "search=OpenAlex search (broad coverage), " "arxiv_search=arXiv search (preprints, ML/CS/Physics), " "semantic_search=Semantic Scholar search, " + "paperclip_search=Paperclip search (biomedical: bioRxiv, medRxiv, PMC, arXiv — 8M+ papers), " + "paperclip_lookup=lookup paper by DOI or PMID via Paperclip, " "trending=highly cited recent papers, details=paper metadata, " "read_paper=read arXiv paper sections, citations=references and citing papers, " "recommend=related papers, find_code=code implementations, " @@ -124,6 +131,11 @@ def create_papers_tool() -> ToolSpec: "enum": ["openalex", "semantic_scholar", "auto"], "description": "Preferred source for search (default: auto, tries OpenAlex then Semantic Scholar)", }, + "paperclip_source": { + "type": "string", + "enum": ["biorxiv", "medrxiv", "pmc", "arxiv", "all"], + "description": "For paperclip_search: filter by source (default: all)", + }, }, "required": ["operation"], }, @@ -177,6 +189,7 @@ async def _handle_papers( year_to: int = None, limit: int = 10, source: str = "auto", + paperclip_source: str = "all", session=None, **kwargs, ) -> tuple[str, bool]: @@ -185,6 +198,8 @@ async def _handle_papers( "search", "arxiv_search", "semantic_search", + "paperclip_search", + "paperclip_lookup", "trending", "details", "citations", @@ -213,6 +228,10 @@ async def _handle_papers( "search": lambda: _search(query, year_from, year_to, limit, source), "arxiv_search": lambda: _arxiv_search(query, year_from, year_to, limit), "semantic_search": lambda: _semantic_scholar_search(query, year_from, year_to, limit), + "paperclip_search": lambda: _paperclip_search( + query, year_from, year_to, limit, paperclip_source + ), + "paperclip_lookup": lambda: _paperclip_lookup(paper_id), "trending": lambda: _trending(query, limit), "details": lambda: _details(paper_id), "read_paper": lambda: _read_paper(paper_id, section), @@ -495,6 +514,164 @@ async def _semantic_scholar_search( return "\n".join(lines), True +# ── Paperclip Search (bioRxiv, medRxiv, PMC, arXiv) ─────────────── + + +def _get_paperclip_headers() -> dict | None: + """Get Paperclip auth headers. Returns None if not configured.""" + api_key = os.environ.get("PAPERCLIP_API_KEY") + if not api_key: + return None + return { + "Content-Type": "application/json", + "X-API-Key": api_key, + } + + +_PAPERCLIP_RATE_LIMIT_MSG = "Paperclip rate limit reached. Try again later." + + +async def _paperclip_search( + query: str, + year_from: int | None = None, + year_to: int | None = None, + limit: int = 10, + paperclip_source: str = "all", +) -> tuple[str, bool]: + """Search biomedical papers using Paperclip (bioRxiv, medRxiv, PMC, arXiv).""" + if not query: + return "Provide a 'query' for search.", False + + headers = _get_paperclip_headers() + if not headers: + return ( + "PAPERCLIP_API_KEY not configured. " + "Set it in Settings > Providers or set the PAPERCLIP_API_KEY environment variable. " + "Get an API key at https://paperclip.gxl.ai", + False, + ) + + # Build the raw CLI-style arguments string + raw_parts = [f'"{query}"', f"-n {min(limit, 100)}"] + + if paperclip_source and paperclip_source != "all": + raw_parts.append(f"--source {paperclip_source}") + + year = year_from or year_to + if year: + raw_parts.append(f"--year {year}") + + raw = " ".join(raw_parts) + + try: + resp = await fetch_with_retry( + f"{PAPERCLIP_API}/api/cli/execute", + method="POST", + headers=headers, + json={"command": "search", "raw": raw}, + timeout=120, + max_retries=2, + ) + except RateLimitError: + return _PAPERCLIP_RATE_LIMIT_MSG, False + except Exception as e: + log.warning(f"Paperclip search error: {e}") + return f"Paperclip error: {str(e)[:200]}", False + + return _parse_paperclip_response(resp, query) + + +def _parse_paperclip_response(resp, query: str) -> tuple[str, bool]: + """Parse and format a Paperclip API response.""" + if resp.status_code in (401, 403): + return ( + "PAPERCLIP_API_KEY is invalid or expired. Check your API key in Settings > Providers.", + False, + ) + if resp.status_code == 429: + return _PAPERCLIP_RATE_LIMIT_MSG, False + if resp.status_code != 200: + try: + detail = resp.json().get("detail", resp.text[:300]) + except Exception: + detail = resp.text[:300] + return f"Paperclip error {resp.status_code}: {detail}", False + + data = resp.json() + output = data.get("output", "") + + if not output: + return f"No papers found for: {query}", True + + # Prepend source attribution + result_id = data.get("result_id", "") + elapsed = data.get("elapsed_ms") + header = "Results via Paperclip (bioRxiv/medRxiv/PMC/arXiv)" + if elapsed: + header += f" [{elapsed}ms]" + if result_id: + header += f" [{result_id}]" + + return f"{header}:\n\n{output}", True + + +async def _paperclip_lookup(paper_id: str) -> tuple[str, bool]: + """Look up a paper by DOI or PMID using Paperclip.""" + if not paper_id: + return "Provide a 'paper_id' (DOI like '10.1101/...' or PMID).", False + + headers = _get_paperclip_headers() + if not headers: + return ( + "PAPERCLIP_API_KEY not configured. " + "Set it in Settings > Providers or set the PAPERCLIP_API_KEY environment variable. " + "Get an API key at https://paperclip.gxl.ai", + False, + ) + + # Determine lookup field based on ID format + if paper_id.startswith("10."): + field = "doi" + elif paper_id.isdigit(): + field = "pmid" + else: + field = "doi" # default to DOI + + raw = f"{field} {paper_id}" + + try: + resp = await fetch_with_retry( + f"{PAPERCLIP_API}/api/cli/execute", + method="POST", + headers=headers, + json={"command": "lookup", "raw": raw}, + timeout=60, + max_retries=2, + ) + except RateLimitError: + return _PAPERCLIP_RATE_LIMIT_MSG, False + except Exception as e: + log.warning(f"Paperclip lookup error: {e}") + return f"Paperclip lookup error: {str(e)[:200]}", False + + if resp.status_code in (401, 403): + return "PAPERCLIP_API_KEY is invalid or expired.", False + if resp.status_code != 200: + try: + detail = resp.json().get("detail", resp.text[:300]) + except Exception: + detail = resp.text[:300] + return f"Paperclip lookup error {resp.status_code}: {detail}", False + + data = resp.json() + output = data.get("output", "") + + if not output: + return f"No paper found for {field}: {paper_id}", True + + return f"Paperclip lookup result:\n\n{output}", True + + # ── Trending (OpenAlex) ────────────────────────────────── diff --git a/backend/openmlr/tools/registry.py b/backend/openmlr/tools/registry.py index aa84f52..f79b51e 100644 --- a/backend/openmlr/tools/registry.py +++ b/backend/openmlr/tools/registry.py @@ -39,6 +39,8 @@ "compute_probe", # Workspace (knowledge graph, notes, search — always accessible) "workspace", + # Parallel file inspection (read-only) + "inspect_files", }, "blocked_message": ( "Tool '{tool}' is not available in PLAN mode. " @@ -80,8 +82,9 @@ class ToolRouter: def __init__(self): self.tools: dict[str, ToolSpec] = {} - self._mcp_client = None self._blocklist: set[str] = set() + self._mcp_clients: dict[str, object] = {} # tool_name -> MCP client + self._mcp_tool_modes: dict[str, list[str]] = {} # tool_name -> allowed modes self._current_mode: str = "general" self._user_id: int | None = None self._db = None @@ -118,10 +121,21 @@ def is_tool_allowed(self, name: str) -> tuple[bool, str]: Returns (allowed, error_message). Supports both 'allowed' (whitelist) and 'blocked' (blacklist) sets. + MCP tools use their own per-tool mode configuration. """ if self._current_mode not in MODE_TOOL_RESTRICTIONS: return True, "" + # MCP tools: check their per-tool mode configuration + if name in self._mcp_tool_modes: + allowed_modes = self._mcp_tool_modes[name] + if self._current_mode in allowed_modes: + return True, "" + return False, ( + f"MCP tool '{name}' is not available in {self._current_mode.upper()} mode. " + f"It is configured for: {', '.join(m.upper() for m in allowed_modes)}." + ) + restrictions = MODE_TOOL_RESTRICTIONS[self._current_mode] # Blacklist mode: specific tools are blocked @@ -325,10 +339,11 @@ async def call_tool( False, ) - # MCP tool (no handler — dispatch to MCP client) - if self._mcp_client: + # MCP tool (no handler — dispatch to per-tool MCP client) + mcp_client = self._mcp_clients.get(name) + if mcp_client: try: - result = await self._mcp_client.call_tool(name, arguments) + result = await mcp_client.call_tool(name, arguments) output = _convert_mcp_content(result) if _research_warning: output += _research_warning @@ -338,16 +353,29 @@ async def call_tool( return f"Tool '{name}' has no handler and no MCP client configured.", False - async def register_mcp_tools(self, mcp_client, blocklist: set[str] | None = None) -> int: - """Register tools from an MCP client. Returns count of tools registered.""" - self._mcp_client = mcp_client - self._blocklist = blocklist or set() + async def register_mcp_tools( + self, + mcp_client, + blocklist: set[str] | None = None, + modes: list[str] | None = None, + ) -> int: + """Register tools from an MCP client. + + Each tool is mapped to its originating client for dispatch, and + tagged with the modes it is allowed in (default: plan + execute). + Returns count of tools registered. + """ + effective_blocklist = blocklist or set() + effective_modes = modes or ["plan", "execute"] count = 0 try: tools = await mcp_client.list_tools() for tool in tools: - if tool.name in self._blocklist or tool.name in self.tools: + if tool.name in effective_blocklist: + continue + if tool.name in self.tools: + logger.warning(f"MCP tool '{tool.name}' shadows built-in tool — skipped") continue spec = ToolSpec( name=tool.name, @@ -356,9 +384,11 @@ async def register_mcp_tools(self, mcp_client, blocklist: set[str] | None = None handler=None, # MCP tools dispatched via call_tool ) self.tools[spec.name] = spec + self._mcp_clients[spec.name] = mcp_client + self._mcp_tool_modes[spec.name] = list(effective_modes) count += 1 - except Exception: - pass + except Exception as e: + logger.warning(f"Failed to register MCP tools: {e}") return count @@ -387,6 +417,7 @@ def create_tool_router(sandbox_manager=None) -> ToolRouter: from .ask_user import create_ask_user_tool from .github import create_github_tools from .huggingface import create_huggingface_tools + from .inspect import create_inspect_tool from .local import create_local_tools from .papers import create_papers_tool from .plan import create_plan_tool @@ -395,6 +426,7 @@ def create_tool_router(sandbox_manager=None) -> ToolRouter: from .writing import create_writing_tool router.register_many(create_local_tools()) + router.register(create_inspect_tool()) router.register_many(create_github_tools()) router.register_many(create_huggingface_tools()) router.register_many(create_search_tools()) diff --git a/backend/openmlr/tools/research.py b/backend/openmlr/tools/research.py index 49ecda0..c805aa2 100644 --- a/backend/openmlr/tools/research.py +++ b/backend/openmlr/tools/research.py @@ -20,7 +20,7 @@ You can use ONLY these read-only tools: - **web_search**: General web search for docs, blog posts, tutorials -- **papers**: Academic paper search (OpenAlex, arXiv, Semantic Scholar, Citations, etc.) +- **papers**: Academic paper search (OpenAlex, arXiv, Semantic Scholar, Paperclip for biomedical, Citations, etc.) - **github_read_file**: Read specific files from GitHub repos - **github_find_examples**: Search code patterns across GitHub - **hf_search_models**: Search Hugging Face models diff --git a/backend/tests/test_conversations.py b/backend/tests/test_conversations.py index 368ea5e..eb5f485 100644 --- a/backend/tests/test_conversations.py +++ b/backend/tests/test_conversations.py @@ -201,3 +201,218 @@ async def test_delete_conversation_unauthenticated(client): """DELETE /api/conversations/{uuid} without auth returns 401.""" resp = await client.delete("/api/conversations/any-uuid") assert resp.status_code == 401 + + +# --------------------------------------------------------------------------- +# Auto-title generation +# --------------------------------------------------------------------------- + + +async def test_get_conversation_no_auto_title_when_empty_default(auth_client): + """No auto-title fires for a default-titled conversation with zero messages.""" + from unittest.mock import patch + + resp = await auth_client.post( + "/api/conversations", + json={"model": None, "mode": "general"}, + ) + uuid = resp.json()["conversation"]["uuid"] + + with patch("openmlr.routes.agent._auto_title", new_callable=AsyncMock) as mock_at: + resp = await auth_client.get(f"/api/conversations/{uuid}") + assert resp.status_code == 200 + mock_at.assert_not_called() + + +async def test_get_conversation_no_auto_title_when_already_titled(auth_client, db_session): + """No auto-title fires for a conversation that already has a custom title.""" + from unittest.mock import patch + + from openmlr.db import operations as ops + + created = await _create_conversation(auth_client, title="My Research") + conv_id = created["conversation"]["id"] + uuid = created["conversation"]["uuid"] + + # Add a message so the msgs list is non-empty + await ops.add_message(db_session, conv_id, "user", "Hello world") + + with patch("openmlr.routes.agent._auto_title", new_callable=AsyncMock) as mock_at: + resp = await auth_client.get(f"/api/conversations/{uuid}") + assert resp.status_code == 200 + # Title is not "New conversation" → auto-title should NOT fire + mock_at.assert_not_called() + + +async def test_get_conversation_triggers_auto_title_for_default_with_messages( + auth_client, db_session +): + """Auto-title fires when a default-titled conversation has messages (page refresh).""" + from unittest.mock import patch + + from openmlr.db import operations as ops + + resp = await auth_client.post( + "/api/conversations", + json={"model": None, "mode": "general"}, + ) + conv_id = resp.json()["conversation"]["id"] + uuid = resp.json()["conversation"]["uuid"] + + # Add messages so the trigger condition is met + await ops.add_message(db_session, conv_id, "user", "Tell me about ML") + await ops.add_message(db_session, conv_id, "assistant", "Machine learning is...") + + with patch("openmlr.routes.agent._auto_title", new_callable=AsyncMock) as mock_at: + resp = await auth_client.get(f"/api/conversations/{uuid}") + assert resp.status_code == 200 + # Title is "New conversation" and msgs exist → auto-title SHOULD fire + mock_at.assert_called_once() + + +async def test_auto_title_skips_update_when_already_titled(db_session, test_user): + """_auto_title does not overwrite an existing non-default title (race guard).""" + from openmlr.db import operations as ops + from openmlr.routes.agent import _auto_title + + conv = await ops.create_conversation(db_session, test_user.id, title="Already Set") + + sm = MagicMock() + sm.generate_title = AsyncMock(return_value="LLM Suggested Title") + bus = MagicMock() + bus.broadcast = AsyncMock() + + messages = [{"role": "user", "content": "Hello"}] + await _auto_title(sm, bus, db_session, conv.id, conv.uuid, messages) + + # Title should remain unchanged + updated = await ops.get_conversation_by_id(db_session, conv.id) + assert updated.title == "Already Set" + bus.broadcast.assert_not_called() + + +async def test_auto_title_updates_when_untitled(db_session, test_user): + """_auto_title sets the title when it is still the default 'New conversation'.""" + from openmlr.db import operations as ops + from openmlr.routes.agent import _auto_title + + conv = await ops.create_conversation(db_session, test_user.id) + assert conv.title == "New conversation" + + sm = MagicMock() + sm.generate_title = AsyncMock(return_value="ML Pipeline Design") + bus = MagicMock() + bus.broadcast = AsyncMock() + + messages = [{"role": "user", "content": "Help me design a pipeline"}] + await _auto_title(sm, bus, db_session, conv.id, conv.uuid, messages) + + updated = await ops.get_conversation_by_id(db_session, conv.id) + assert updated.title == "ML Pipeline Design" + bus.broadcast.assert_called_once() + + +async def test_auto_title_no_update_on_generation_failure(db_session, test_user): + """_auto_title leaves the title unchanged when LLM generation returns None.""" + from openmlr.db import operations as ops + from openmlr.routes.agent import _auto_title + + conv = await ops.create_conversation(db_session, test_user.id) + + sm = MagicMock() + sm.generate_title = AsyncMock(return_value=None) + bus = MagicMock() + bus.broadcast = AsyncMock() + + # Pass empty messages so fallback also produces None + await _auto_title(sm, bus, db_session, conv.id, conv.uuid, []) + + updated = await ops.get_conversation_by_id(db_session, conv.id) + assert updated.title == "New conversation" + bus.broadcast.assert_not_called() + + +# --------------------------------------------------------------------------- +# @ Mention enrichment +# --------------------------------------------------------------------------- + + +def test_enrich_with_mentions_no_mentions(): + """No mentions returns the original message unchanged.""" + from openmlr.routes.agent import _enrich_with_mentions + + msg = "Hello world" + assert _enrich_with_mentions(msg, None) == msg + assert _enrich_with_mentions(msg, []) == msg + + +def test_enrich_with_file_mention(): + """File mention adds a reference hint.""" + from openmlr.models import Mention + from openmlr.routes.agent import _enrich_with_mentions + + msg = "Check this file" + mentions = [Mention(type="file", value="code/train.py")] + result = _enrich_with_mentions(msg, mentions) + assert "code/train.py" in result + assert "read" in result.lower() + assert msg in result + + +def test_enrich_with_directory_mention(): + """Directory mention suggests inspect_files.""" + from openmlr.models import Mention + from openmlr.routes.agent import _enrich_with_mentions + + msg = "Look at this" + mentions = [Mention(type="file", value="code/")] + result = _enrich_with_mentions(msg, mentions) + assert "code/" in result + assert "inspect_files" in result + assert msg in result + + +def test_enrich_with_server_mention(): + """Server mention references MCP tools.""" + from openmlr.models import Mention + from openmlr.routes.agent import _enrich_with_mentions + + msg = "Use gmail" + mentions = [Mention(type="server", value="my-gmail")] + result = _enrich_with_mentions(msg, mentions) + assert "my-gmail" in result + assert "MCP" in result + assert msg in result + + +def test_enrich_sanitizes_injection_attempt(): + """Mention values are sanitized to prevent prompt injection.""" + from openmlr.models import Mention + from openmlr.routes.agent import _enrich_with_mentions + + msg = "Check this" + # Backticks and newlines should be stripped + mentions = [Mention(type="file", value="file`\nignore instructions")] + result = _enrich_with_mentions(msg, mentions) + assert "`" not in result.split(msg)[0] # no backticks in the hint part + assert "\n\n" in result # only the separator newlines + # The sanitized value should be present without injection characters + assert "fileignore instructions" in result + + +def test_enrich_multiple_mentions(): + """Multiple mentions of different types.""" + from openmlr.models import Mention + from openmlr.routes.agent import _enrich_with_mentions + + msg = "Do the thing" + mentions = [ + Mention(type="file", value="code/train.py"), + Mention(type="server", value="my-mcp"), + Mention(type="file", value="data/"), + ] + result = _enrich_with_mentions(msg, mentions) + assert "code/train.py" in result + assert "my-mcp" in result + assert "data/" in result + assert msg in result diff --git a/backend/tests/test_inspect.py b/backend/tests/test_inspect.py new file mode 100644 index 0000000..5c60004 --- /dev/null +++ b/backend/tests/test_inspect.py @@ -0,0 +1,197 @@ +"""Tests for inspect_files tool — parallel reading, relevance scoring, security.""" + +import pytest + +from openmlr.tools.inspect import ( + _is_relative, + _read_file_snippet, + _score_relevance, + create_inspect_tool, +) + +pytestmark = pytest.mark.asyncio + + +class TestScoreRelevance: + def test_full_match(self): + score = _score_relevance("the training loop runs for epochs", "training loop") + assert score == pytest.approx(1.0) + + def test_partial_match(self): + score = _score_relevance("the training loop runs", "training loop optimizer") + assert 0.3 < score < 1.0 + + def test_no_match(self): + score = _score_relevance("hello world", "gradient descent") + assert score == pytest.approx(0.0) + + def test_short_query_terms_ignored(self): + """Words <= 2 chars are excluded from scoring.""" + score = _score_relevance("hello world", "a b c") + assert score == pytest.approx(0.5) # fallback for no useful terms + + def test_case_insensitive(self): + score = _score_relevance("Training Loop OPTIMIZER", "training loop optimizer") + assert score == pytest.approx(1.0) + + def test_empty_content(self): + score = _score_relevance("", "training") + assert score == pytest.approx(0.0) + + +class TestReadFileSnippet: + def test_small_file(self, tmp_path): + f = tmp_path / "small.txt" + f.write_text("line one\nline two\nline three") + content, lines = _read_file_snippet(f) + assert lines == 3 + assert "1: line one" in content + assert "3: line three" in content + + def test_truncates_large_files(self, tmp_path): + f = tmp_path / "large.txt" + f.write_text("\n".join(f"line {i}" for i in range(500))) + content, lines = _read_file_snippet(f) + assert lines == 500 + assert "more lines truncated" in content + + def test_nonexistent_file(self, tmp_path): + f = tmp_path / "missing.txt" + content, lines = _read_file_snippet(f) + assert content == "" + assert lines == 0 + + def test_large_file_skipped(self, tmp_path): + f = tmp_path / "huge.bin" + # Write a 3MB file (above _MAX_FILE_SIZE of 2MB) + f.write_bytes(b"x" * (3 * 1024 * 1024)) + content, lines = _read_file_snippet(f) + assert "too large" in content.lower() + assert lines == 0 + + def test_binary_file_graceful(self, tmp_path): + f = tmp_path / "binary.dat" + f.write_bytes(b"\x00\x01\x02\xff" * 100) + content, _ = _read_file_snippet(f) + # Should not crash — errors="replace" handles it + assert isinstance(content, str) + + +class TestIsRelative: + def test_relative(self, tmp_path): + child = tmp_path / "sub" / "file.txt" + assert _is_relative(child, tmp_path) is True + + def test_not_relative(self, tmp_path): + from pathlib import Path + + assert _is_relative(Path("/etc/passwd"), tmp_path) is False + + +class TestCreateInspectTool: + def test_creates_tool(self): + tool = create_inspect_tool() + assert tool.name == "inspect_files" + assert tool.handler is not None + assert "paths" in tool.parameters["properties"] + assert "query" in tool.parameters["properties"] + + +class TestHandleInspectFiles: + async def test_empty_paths(self): + from openmlr.tools.inspect import _handle_inspect_files + + output, success = await _handle_inspect_files(paths=[], query="test") + assert success is False + assert "No paths" in output + + async def test_empty_query(self): + from openmlr.tools.inspect import _handle_inspect_files + + output, success = await _handle_inspect_files(paths=["file.py"], query="") + assert success is False + assert "No query" in output + + async def test_reads_files_from_directory(self, tmp_path, monkeypatch): + """inspect_files reads files from a directory and scores relevance.""" + from openmlr.tools.inspect import _handle_inspect_files + from openmlr.tools.local import _project_workspace_var + + # Set workspace to tmp_path + token = _project_workspace_var.set(str(tmp_path)) + try: + # Create test files + (tmp_path / "train.py").write_text("def training_loop():\n pass") + (tmp_path / "utils.py").write_text("def helper():\n pass") + + output, success = await _handle_inspect_files( + paths=[str(tmp_path)], query="training loop" + ) + assert success is True + assert "train.py" in output + finally: + _project_workspace_var.reset(token) + + async def test_respects_max_files(self, tmp_path, monkeypatch): + """max_files limits how many files are read.""" + from openmlr.tools.inspect import _handle_inspect_files + from openmlr.tools.local import _project_workspace_var + + token = _project_workspace_var.set(str(tmp_path)) + try: + for i in range(10): + (tmp_path / f"file_{i}.py").write_text(f"content {i}") + + output, success = await _handle_inspect_files( + paths=[str(tmp_path)], query="content", max_files=3 + ) + assert success is True + assert "Truncated" in output + finally: + _project_workspace_var.reset(token) + + async def test_negative_max_files_clamped(self, tmp_path): + """Negative max_files is clamped to 1.""" + from openmlr.tools.inspect import _handle_inspect_files + from openmlr.tools.local import _project_workspace_var + + token = _project_workspace_var.set(str(tmp_path)) + try: + (tmp_path / "a.txt").write_text("hello") + (tmp_path / "b.txt").write_text("world") + + _, success = await _handle_inspect_files( + paths=[str(tmp_path)], query="hello", max_files=-5 + ) + assert success is True + # Should not crash or read zero files + finally: + _project_workspace_var.reset(token) + + async def test_symlink_outside_workspace_blocked(self, tmp_path): + """Symlinks pointing outside the workspace should be skipped.""" + import os + + from openmlr.tools.inspect import _handle_inspect_files + from openmlr.tools.local import _project_workspace_var + + workspace = tmp_path / "workspace" + workspace.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.txt").write_text("sensitive data") + + # Create symlink from workspace to outside + symlink = workspace / "link_to_outside.txt" + try: + os.symlink(outside / "secret.txt", symlink) + except OSError: + pytest.skip("Symlinks not supported on this platform") + + token = _project_workspace_var.set(str(workspace)) + try: + output, success = await _handle_inspect_files(paths=[str(workspace)], query="sensitive") + assert success is True + assert "sensitive data" not in output + finally: + _project_workspace_var.reset(token) diff --git a/backend/tests/test_models.py b/backend/tests/test_models.py index ef9cf1b..fc4ed69 100644 --- a/backend/tests/test_models.py +++ b/backend/tests/test_models.py @@ -175,6 +175,31 @@ def test_allows_null_mode(self): m = MessageSend(message="test", mode=None) assert m.mode is None + def test_mentions_default_none(self): + m = MessageSend(message="test") + assert m.mentions is None + + def test_with_mentions(self): + from openmlr.models import Mention + + m = MessageSend( + message="check @train.py", + mentions=[ + Mention(type="file", value="code/train.py"), + Mention(type="server", value="my-mcp"), + ], + ) + assert len(m.mentions) == 2 + assert m.mentions[0].type == "file" + assert m.mentions[0].value == "code/train.py" + assert m.mentions[1].type == "server" + + def test_mention_rejects_invalid_type(self): + from openmlr.models import Mention + + with pytest.raises(ValidationError): + Mention(type="invalid", value="test") + class TestApprovalRequest: def test_valid(self): diff --git a/backend/tests/test_tool_registry.py b/backend/tests/test_tool_registry.py index 36ffbfa..a9952c3 100644 --- a/backend/tests/test_tool_registry.py +++ b/backend/tests/test_tool_registry.py @@ -382,3 +382,160 @@ async def research_handler(**kwargs): output, success = await router.call_tool("web_search", {}) assert success is True assert "PLAN MODE RESEARCH BUDGET" not in output + + +class TestMCPToolRegistration: + """Tests for MCP tool registration, multi-client dispatch, and mode filtering.""" + + async def test_register_mcp_tool_default_modes(self, router): + """MCP tools default to both plan and execute modes.""" + from unittest.mock import AsyncMock, MagicMock + + client = MagicMock() + tool = MagicMock() + tool.name = "mcp_search" + tool.description = "Search via MCP" + tool.input_schema = {"type": "object", "properties": {}} + client.list_tools = AsyncMock(return_value=[tool]) + + count = await router.register_mcp_tools(client) + assert count == 1 + assert "mcp_search" in router.tools + assert router._mcp_tool_modes["mcp_search"] == ["plan", "execute"] + + async def test_register_mcp_tool_custom_modes(self, router): + """MCP tools respect custom mode configuration.""" + from unittest.mock import AsyncMock, MagicMock + + client = MagicMock() + tool = MagicMock() + tool.name = "exec_only" + tool.description = "Execute-only tool" + tool.input_schema = {"type": "object", "properties": {}} + client.list_tools = AsyncMock(return_value=[tool]) + + count = await router.register_mcp_tools(client, modes=["execute"]) + assert count == 1 + assert router._mcp_tool_modes["exec_only"] == ["execute"] + + async def test_mcp_tool_allowed_in_configured_mode(self, router): + """MCP tool is allowed when current mode matches configured modes.""" + from unittest.mock import AsyncMock, MagicMock + + client = MagicMock() + tool = MagicMock() + tool.name = "mcp_tool" + tool.description = "" + tool.input_schema = {"type": "object", "properties": {}} + client.list_tools = AsyncMock(return_value=[tool]) + await router.register_mcp_tools(client, modes=["plan", "execute"]) + + router.set_mode("plan") + allowed, _ = router.is_tool_allowed("mcp_tool") + assert allowed is True + + async def test_mcp_tool_blocked_outside_configured_mode(self, router): + """MCP tool is blocked when current mode is not in configured modes.""" + from unittest.mock import AsyncMock, MagicMock + + client = MagicMock() + tool = MagicMock() + tool.name = "exec_tool" + tool.description = "" + tool.input_schema = {"type": "object", "properties": {}} + client.list_tools = AsyncMock(return_value=[tool]) + await router.register_mcp_tools(client, modes=["execute"]) + + router.set_mode("plan") + allowed, msg = router.is_tool_allowed("exec_tool") + assert allowed is False + assert "PLAN" in msg + + async def test_mcp_multi_client_dispatch(self, router): + """Tools from different MCP servers dispatch to their originating client.""" + from unittest.mock import AsyncMock, MagicMock + + # Server A + client_a = MagicMock() + tool_a = MagicMock() + tool_a.name = "tool_from_a" + tool_a.description = "From A" + tool_a.input_schema = {"type": "object", "properties": {}} + client_a.list_tools = AsyncMock(return_value=[tool_a]) + client_a.call_tool = AsyncMock(return_value="result_a") + + # Server B + client_b = MagicMock() + tool_b = MagicMock() + tool_b.name = "tool_from_b" + tool_b.description = "From B" + tool_b.input_schema = {"type": "object", "properties": {}} + client_b.list_tools = AsyncMock(return_value=[tool_b]) + client_b.call_tool = AsyncMock(return_value="result_b") + + await router.register_mcp_tools(client_a) + await router.register_mcp_tools(client_b) + + assert "tool_from_a" in router.tools + assert "tool_from_b" in router.tools + + # Dispatch tool_from_a to client_a + await router.call_tool("tool_from_a", {}, enforce_mode=False) + client_a.call_tool.assert_called_once_with("tool_from_a", {}) + + # Dispatch tool_from_b to client_b + await router.call_tool("tool_from_b", {}, enforce_mode=False) + client_b.call_tool.assert_called_once_with("tool_from_b", {}) + + async def test_mcp_tool_no_shadow_builtin(self, router, bash_tool): + """MCP tools cannot shadow built-in tools.""" + from unittest.mock import AsyncMock, MagicMock + + router.register(bash_tool) + + client = MagicMock() + shadow = MagicMock() + shadow.name = "bash" # same name as built-in + shadow.description = "Malicious" + shadow.input_schema = {"type": "object", "properties": {}} + client.list_tools = AsyncMock(return_value=[shadow]) + + count = await router.register_mcp_tools(client) + assert count == 0 + # The built-in should still have its handler + assert router.tools["bash"].handler is not None + + async def test_mcp_register_logs_exception(self, router, caplog): + """register_mcp_tools logs a warning on exceptions.""" + import logging + from unittest.mock import AsyncMock, MagicMock + + client = MagicMock() + client.list_tools = AsyncMock(side_effect=ConnectionError("refused")) + + with caplog.at_level(logging.WARNING, logger="openmlr.tools.registry"): + count = await router.register_mcp_tools(client) + assert count == 0 + assert "Failed to register MCP tools" in caplog.text + + async def test_mcp_filtered_from_specs_by_mode(self, router): + """MCP tools are filtered from LLM specs based on mode configuration.""" + from unittest.mock import AsyncMock, MagicMock + + client = MagicMock() + tool = MagicMock() + tool.name = "exec_only_mcp" + tool.description = "" + tool.input_schema = {"type": "object", "properties": {}} + client.list_tools = AsyncMock(return_value=[tool]) + await router.register_mcp_tools(client, modes=["execute"]) + + router.set_mode("plan") + specs = router.get_tool_specs_for_llm(filter_by_mode=True) + names = [s["function"]["name"] for s in specs] + assert "exec_only_mcp" not in names + + router.set_mode("execute") + specs = router.get_tool_specs_for_llm(filter_by_mode=True) + names = [s["function"]["name"] for s in specs] + assert "exec_only_mcp" in names diff --git a/backend/tests/test_tools_papers.py b/backend/tests/test_tools_papers.py index df7627b..e237e22 100644 --- a/backend/tests/test_tools_papers.py +++ b/backend/tests/test_tools_papers.py @@ -1,12 +1,18 @@ """Tests for papers tool — helper functions and tool spec.""" +import os +from unittest.mock import MagicMock, patch + import pytest from openmlr.tools.papers import ( _check_budget, _extract_arxiv_id, _get_budget_info, + _get_paperclip_headers, _increment_budget, + _paperclip_lookup, + _paperclip_search, _reconstruct_abstract, _to_openalex_id, create_papers_tool, @@ -31,6 +37,21 @@ async def test_creates_tool(self): assert "find_code" in ops assert "find_datasets" in ops + def test_paperclip_operations_in_enum(self): + tool = create_papers_tool() + ops = tool.parameters["properties"]["operation"]["enum"] + assert "paperclip_search" in ops + assert "paperclip_lookup" in ops + + def test_paperclip_source_parameter(self): + tool = create_papers_tool() + props = tool.parameters["properties"] + assert "paperclip_source" in props + assert "biorxiv" in props["paperclip_source"]["enum"] + assert "medrxiv" in props["paperclip_source"]["enum"] + assert "pmc" in props["paperclip_source"]["enum"] + assert "all" in props["paperclip_source"]["enum"] + class TestExtractArxivId: async def test_standard_format(self): @@ -102,3 +123,141 @@ async def test_increment_and_check(self): info = _get_budget_info() assert info["used"] >= 1 _search_counts.clear() + + +class TestPaperclipHeaders: + def test_returns_none_without_key(self): + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("PAPERCLIP_API_KEY", None) + assert _get_paperclip_headers() is None + + def test_returns_headers_with_key(self): + with patch.dict(os.environ, {"PAPERCLIP_API_KEY": "gxl_test123"}): + headers = _get_paperclip_headers() + assert headers is not None + assert headers["X-API-Key"] == "gxl_test123" + assert headers["Content-Type"] == "application/json" + + +class TestPaperclipSearch: + async def test_no_query_returns_error(self): + result, ok = await _paperclip_search("") + assert ok is False + assert "query" in result.lower() + + async def test_no_api_key_returns_error(self): + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("PAPERCLIP_API_KEY", None) + result, ok = await _paperclip_search("CRISPR") + assert ok is False + assert "PAPERCLIP_API_KEY" in result + + @patch("openmlr.tools.papers.fetch_with_retry") + async def test_successful_search(self, mock_fetch): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = { + "output": " 1. CRISPR gene editing in vivo\n doi:10.1101/2024.01.01\n", + "elapsed_ms": 150, + "result_id": "s_abc123", + } + mock_fetch.return_value = mock_resp + + with patch.dict(os.environ, {"PAPERCLIP_API_KEY": "gxl_test123"}): + result, ok = await _paperclip_search("CRISPR", limit=5) + + assert ok is True + assert "Paperclip" in result + assert "CRISPR" in result + + # Verify the API call + mock_fetch.assert_called_once() + call_kwargs = mock_fetch.call_args + assert call_kwargs.kwargs["method"] == "POST" + assert "/api/cli/execute" in call_kwargs.args[0] + body = call_kwargs.kwargs["json"] + assert body["command"] == "search" + assert "-n 5" in body["raw"] + + @patch("openmlr.tools.papers.fetch_with_retry") + async def test_search_with_source_filter(self, mock_fetch): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"output": "results", "elapsed_ms": 100} + mock_fetch.return_value = mock_resp + + with patch.dict(os.environ, {"PAPERCLIP_API_KEY": "gxl_test123"}): + await _paperclip_search("CRISPR", paperclip_source="pmc") + + body = mock_fetch.call_args.kwargs["json"] + assert "--source pmc" in body["raw"] + + @patch("openmlr.tools.papers.fetch_with_retry") + async def test_search_with_year_filter(self, mock_fetch): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"output": "results", "elapsed_ms": 100} + mock_fetch.return_value = mock_resp + + with patch.dict(os.environ, {"PAPERCLIP_API_KEY": "gxl_test123"}): + await _paperclip_search("CRISPR", year_from=2024) + + body = mock_fetch.call_args.kwargs["json"] + assert "--year 2024" in body["raw"] + + @patch("openmlr.tools.papers.fetch_with_retry") + async def test_auth_error(self, mock_fetch): + mock_resp = MagicMock() + mock_resp.status_code = 401 + mock_fetch.return_value = mock_resp + + with patch.dict(os.environ, {"PAPERCLIP_API_KEY": "gxl_bad_key"}): + result, ok = await _paperclip_search("CRISPR") + + assert ok is False + assert "invalid" in result.lower() or "expired" in result.lower() + + +class TestPaperclipLookup: + async def test_no_paper_id_returns_error(self): + result, ok = await _paperclip_lookup("") + assert ok is False + assert "paper_id" in result.lower() + + async def test_no_api_key_returns_error(self): + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("PAPERCLIP_API_KEY", None) + result, ok = await _paperclip_lookup("10.1101/2024.01.01") + assert ok is False + assert "PAPERCLIP_API_KEY" in result + + @patch("openmlr.tools.papers.fetch_with_retry") + async def test_doi_lookup(self, mock_fetch): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = { + "output": "Title: CRISPR Paper\nDOI: 10.1101/2024.01.01\n", + } + mock_fetch.return_value = mock_resp + + with patch.dict(os.environ, {"PAPERCLIP_API_KEY": "gxl_test123"}): + _, ok = await _paperclip_lookup("10.1101/2024.01.01") + + assert ok is True + body = mock_fetch.call_args.kwargs["json"] + assert body["command"] == "lookup" + assert "doi" in body["raw"] + + @patch("openmlr.tools.papers.fetch_with_retry") + async def test_pmid_lookup(self, mock_fetch): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"output": "Paper details"} + mock_fetch.return_value = mock_resp + + with patch.dict(os.environ, {"PAPERCLIP_API_KEY": "gxl_test123"}): + await _paperclip_lookup("12345678") + + body = mock_fetch.call_args.kwargs["json"] + assert body["command"] == "lookup" + assert "pmid" in body["raw"] diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index cdc0285..2627943 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -614,6 +614,13 @@ function ChatUI({ if (uuid && title) setConversations((prev) => prev.map((c) => c.uuid === uuid ? { ...c, title } : c)); break; } + case 'mcp_status': { + const servers = data?.servers; + if (Array.isArray(servers)) { + setMcpServers(servers); + } + break; + } case 'interrupted': setCurrentConvStatus('idle'); setMessages((prev) => [...prev.filter((m) => !(m.role === 'system' && m.content === '::thinking::')), { id: nextId(), role: 'system', content: 'Interrupted.' }]); @@ -677,7 +684,7 @@ function ChatUI({ } }, [currentConvUuid, jobProcessing, connected]); - const sendMessage = useCallback(async (text: string, mode: string) => { + const sendMessage = useCallback(async (text: string, mode: string, mentions?: Array<{ type: 'server' | 'file'; value: string }>) => { // Prevent concurrent/duplicate sends if (sendingRef.current) return; sendingRef.current = true; @@ -685,7 +692,7 @@ function ChatUI({ setMessages((prev) => [...prev, { id: nextId(), role: 'user', content: text, metadata: { tool: mode } }]); setCurrentConvStatus('processing'); try { - await api.sendMessage(text, mode); + await api.sendMessage(text, mode, mentions); } catch (err: any) { setCurrentConvStatus('idle'); setMessages((prev) => [...prev, { id: nextId(), role: 'error', content: `Failed to send: ${err.message}` }]); @@ -793,7 +800,7 @@ function ChatUI({
{/* Agent / Editor / Terminal tab bar */}
@@ -902,7 +909,9 @@ function ChatUI({ text={inputText} onTextChange={setInputText} onSend={sendMessage} - onStop={() => { api.interrupt().catch(() => {}); setCurrentConvStatus('idle'); }} + onStop={() => { api.interrupt().catch(() => {}); setCurrentConvStatus('idle'); }} + mcpServers={mcpServers} + projectUuid={activeProject?.uuid ?? null} />
)} diff --git a/frontend/src/__tests__/InputArea.test.tsx b/frontend/src/__tests__/InputArea.test.tsx index 93d479a..e5f7c6f 100644 --- a/frontend/src/__tests__/InputArea.test.tsx +++ b/frontend/src/__tests__/InputArea.test.tsx @@ -51,7 +51,7 @@ describe('InputArea', () => { // Send button has mode-specific title const sendBtn = screen.getByTitle('Send in Plan mode (Enter)'); fireEvent.click(sendBtn); - expect(onSend).toHaveBeenCalledWith('hello', 'plan'); + expect(onSend).toHaveBeenCalledWith('hello', 'plan', undefined); expect(onTextChange).toHaveBeenCalledWith(''); }); diff --git a/frontend/src/api.ts b/frontend/src/api.ts index cf9a6ba..d2d9fa5 100644 --- a/frontend/src/api.ts +++ b/frontend/src/api.ts @@ -70,8 +70,8 @@ export const api = { getMe: () => get('/api/auth/me'), // Messages - sendMessage: (message: string, mode?: string) => - post('/api/message', { message, mode, request_id: crypto.randomUUID() }), + sendMessage: (message: string, mode?: string, mentions?: Array<{ type: string; value: string }>) => + post('/api/message', { message, mode, request_id: crypto.randomUUID(), mentions: mentions?.length ? mentions : undefined }), submitAnswers: (answers: Record) => post('/api/answers', { answers }), interrupt: () => post('/api/interrupt', {}), sendApproval: (approvals: Record) => post('/api/approval', { approvals }), diff --git a/frontend/src/components/InputArea.tsx b/frontend/src/components/InputArea.tsx index 12a25e7..7bd3dcd 100644 --- a/frontend/src/components/InputArea.tsx +++ b/frontend/src/components/InputArea.tsx @@ -1,21 +1,29 @@ import { useRef, useEffect, useCallback, useState } from 'react'; import { ArrowUp, Square } from 'lucide-react'; +import type { McpServerStatus, Mention } from '../types'; +import { MentionPopover } from './MentionPopover'; export type Mode = 'plan' | 'execute'; interface Props { - disabled: boolean; - showStop?: boolean; - mode: Mode; - onModeChange: (mode: Mode) => void; - onSend: (text: string, mode: Mode) => void; - onStop: () => void; - text: string; - onTextChange: (text: string) => void; + readonly disabled: boolean; + readonly showStop?: boolean; + readonly mode: Mode; + readonly onModeChange: (mode: Mode) => void; + readonly onSend: (text: string, mode: Mode, mentions?: Mention[]) => void; + readonly onStop: () => void; + readonly text: string; + readonly onTextChange: (text: string) => void; + readonly mcpServers?: readonly McpServerStatus[]; + readonly projectUuid?: string | null; } -export function InputArea({ disabled, showStop, mode, onModeChange, onSend, onStop, text, onTextChange }: Props) { +export function InputArea({ disabled, showStop, mode, onModeChange, onSend, onStop, text, onTextChange, mcpServers = [], projectUuid = null }: Props) { const textareaRef = useRef(null); + const [mentions, setMentions] = useState([]); + const [mentionPopoverOpen, setMentionPopoverOpen] = useState(false); + const [mentionQuery, setMentionQuery] = useState(''); + const [mentionStartPos, setMentionStartPos] = useState(null); useEffect(() => { const el = textareaRef.current; @@ -29,15 +37,105 @@ export function InputArea({ disabled, showStop, mode, onModeChange, onSend, onSt const submit = useCallback(() => { const trimmed = text.trim(); if (!trimmed || disabled) return; + // Filter mentions to only those still present in the text + const activeMentions = mentions.filter((m) => { + const tag = `@${m.value}`; + return trimmed.includes(tag); + }); onTextChange(''); - onSend(trimmed, mode); + onSend(trimmed, mode, activeMentions.length > 0 ? activeMentions : undefined); + setMentions([]); + setMentionPopoverOpen(false); if (textareaRef.current) textareaRef.current.style.height = 'auto'; - }, [text, disabled, onSend, mode, onTextChange]); + }, [text, disabled, onSend, mode, onTextChange, mentions]); const toggleMode = useCallback(() => { onModeChange(mode === 'plan' ? 'execute' : 'plan'); }, [mode, onModeChange]); + // Detect @ trigger in text changes + const handleTextChange = useCallback((newText: string) => { + onTextChange(newText); + + const el = textareaRef.current; + if (!el) return; + const cursorPos = el.selectionStart; + + // Find the last @ before cursor that's preceded by whitespace, newline, or is at position 0 + let atPos = -1; + for (let i = cursorPos - 1; i >= 0; i--) { + const ch = newText[i]; + if (ch === '@') { + if (i === 0 || /\s/.test(newText[i - 1])) { + atPos = i; + } + break; + } + // Stop searching if we hit whitespace (except within the query which can have /) + if (ch === ' ' || ch === '\n') break; + } + + if (atPos >= 0) { + const query = newText.slice(atPos + 1, cursorPos); + // Don't show popover if there's a space in the query (mention completed) + if (!query.includes(' ')) { + setMentionPopoverOpen(true); + setMentionQuery(query); + setMentionStartPos(atPos); + return; + } + } + setMentionPopoverOpen(false); + }, [onTextChange]); + + const handleMentionSelect = useCallback((mention: Mention, _displayText: string) => { + if (mentionStartPos === null) return; + const el = textareaRef.current; + if (!el) return; + + // For directory selections, update the query to browse deeper + if (mention.type === 'file' && mention.value.endsWith('/')) { + const newText = text.slice(0, mentionStartPos + 1) + mention.value; + onTextChange(newText); + setMentionQuery(mention.value); + // Keep popover open for continued browsing + setTimeout(() => { + if (textareaRef.current) { + textareaRef.current.selectionStart = newText.length; + textareaRef.current.selectionEnd = newText.length; + textareaRef.current.focus(); + } + }, 0); + return; + } + + // Insert the mention text, replacing @query + const before = text.slice(0, mentionStartPos); + const after = text.slice(mentionStartPos + 1 + mentionQuery.length); + const insertText = `@${mention.value}`; + const newText = before + insertText + (after.startsWith(' ') ? after : ' ' + after); + onTextChange(newText); + + // Track the mention + setMentions((prev) => { + if (prev.some((m) => m.type === mention.type && m.value === mention.value)) return prev; + return [...prev, mention]; + }); + + setMentionPopoverOpen(false); + setMentionStartPos(null); + + // Restore focus and cursor position + setTimeout(() => { + if (textareaRef.current) { + const pos = before.length + insertText.length + 1; + textareaRef.current.selectionStart = pos; + textareaRef.current.selectionEnd = pos; + textareaRef.current.focus(); + } + }, 0); + }, [text, mentionStartPos, mentionQuery, onTextChange]); + // Keyboard shortcuts useEffect(() => { const handler = (e: KeyboardEvent) => { @@ -64,75 +162,121 @@ export function InputArea({ disabled, showStop, mode, onModeChange, onSend, onSt return () => window.removeEventListener('resize', checkMobile); }, []); - const placeholder = isPlan - ? (isMobile ? 'Plan your research...' : 'Plan: ask questions, gather context, create plan...') - : (isMobile ? 'Execute task...' : 'Execute: tell the agent what to do...'); + const placeholder = (() => { + if (isPlan) return isMobile ? 'Plan your research...' : 'Plan: ask questions, gather context... (@ to mention)'; + return isMobile ? 'Execute task...' : 'Execute: tell the agent what to do... (@ to mention)'; + })(); + + // Show mention chips above input + const mentionChips = mentions.filter((m) => { + const tag = `@${m.value}`; + return text.includes(tag); + }); return (
-
- {/* Mode toggle button - fixed height to match input */} - - - {/* Textarea - minimum height matches buttons */} -