From 28a0e5058aedafeca8124f152b7aacd4f68d24bb Mon Sep 17 00:00:00 2001 From: Haoming Yan Date: Thu, 16 Jul 2026 16:48:20 +0800 Subject: [PATCH 1/2] feat(ai): secure providers and complete MCP tools --- CHANGELOG.md | 14 +- MEMORY.md | 3 + README.md | 19 +- docs/agent-memory.md | 13 + package.json | 49 ++- python/openqc/ai/__init__.py | 14 +- python/openqc/ai/client.py | 490 +++++++++++---------- python/openqc/mcp_server.py | 311 +++++++++---- src/ai/AICore.ts | 398 +++++++++-------- src/ai/aiCommands.ts | 67 ++- tests/python/test_ai_client.py | 175 ++++++++ tests/python/test_mcp_contract.py | 112 +++++ tests/unit/ai/AICore.test.ts | 535 +++++++++-------------- tests/unit/ai/aiSecurityContract.test.ts | 20 + 14 files changed, 1374 insertions(+), 846 deletions(-) create mode 100644 MEMORY.md create mode 100644 docs/agent-memory.md create mode 100644 tests/python/test_ai_client.py create mode 100644 tests/python/test_mcp_contract.py create mode 100644 tests/unit/ai/aiSecurityContract.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 1bd933b..d83c43f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,15 @@ +## Unreleased + +### Secure AI providers and MCP + +- Switched the Python AI bridge to a module entry point with the OpenAI Responses API and + Ollama generate API, using fake-HTTP coverage rather than live provider requests. +- Moved OpenAI credentials out of configuration and into VS Code SecretStorage with explicit set + and clear commands. +- Added cancellation, provider and subprocess timeouts, bounded output, and secret-safe errors. +- Routed and validated all five tools declared by the dependency-free MCP server while preserving + JSON-RPC request IDs and normalizing missing optional dependencies and tool failures. + ## [3.0.6] - 2026-03-05 ## [3.0.7] - 2026-03-05 @@ -429,7 +441,7 @@ Phase 5 remaining tasks: - VSCode settings for AI: - `openqc.ai.enabled` - Enable/disable AI features - `openqc.ai.provider` - Select provider (openai/ollama) - - `openqc.ai.apiKey` - OpenAI API key + - OpenAI API key setting (historical; removed above in favor of SecretStorage) - `openqc.ai.model` - Model name - `openqc.ai.ollamaUrl` - Ollama server URL - `openqc.ai.maxTokens` - Maximum response tokens diff --git a/MEMORY.md b/MEMORY.md new file mode 100644 index 0000000..9b2558f --- /dev/null +++ b/MEMORY.md @@ -0,0 +1,3 @@ +# Agent Memory + +- Session notes live in [docs/agent-memory.md](docs/agent-memory.md). diff --git a/README.md b/README.md index d902b8b..3dc9b40 100644 --- a/README.md +++ b/README.md @@ -210,6 +210,21 @@ OpenQC works out of the box, but you can customize it: } ``` +### Secure AI providers + +OpenQC can use the OpenAI Responses API or a local Ollama server for input optimization, +generation, explanation, and debugging. Enable `openqc.ai.enabled`, select the provider and +model, then run **OpenQC: Set OpenAI API Key** when using OpenAI. The credential is stored in +VS Code SecretStorage and is never written to user or workspace settings; use **OpenQC: Clear +OpenAI API Key** to remove it. Provider calls are cancellable and bounded by the configured +timeout, token limit, and output-character limit. + +From a source checkout, the dependency-free MCP entry point is +`PYTHONPATH=python python -m openqc.mcp_server`. It exposes five tools for +structure parsing, output parsing, backend checks, supercell generation, and dataset summaries. +Native paths continue to work when optional scientific packages are absent; tools that require a +missing package return a structured JSON-RPC error without terminating the server. + --- ## 💡 Use Cases @@ -251,8 +266,8 @@ OpenQC works out of the box, but you can customize it: - [ ] Community examples for common calculation workflows ### Long Term (v3.0) -- [ ] AI-powered parameter optimization -- [ ] Natural language input generation +- [x] AI-powered parameter optimization with OpenAI and Ollama +- [x] Natural language input generation - [ ] Workflow automation - [ ] Demand-proven ecosystem integrations from the [roadmap](docs/project/PLAN.md) diff --git a/docs/agent-memory.md b/docs/agent-memory.md new file mode 100644 index 0000000..d280f99 --- /dev/null +++ b/docs/agent-memory.md @@ -0,0 +1,13 @@ +# Agent Memory + +## 2026-07-16 - Issue #237 + +- The production AI bridge must launch `openqc.ai.client` with `python -m`; eager imports from + `openqc.ai.__init__` trigger a `runpy` warning, so client exports are lazy. +- OpenAI uses the Responses API and receives its credential only from VS Code SecretStorage. + Fake local HTTP servers cover provider behavior; tests must never send live provider requests. +- The dependency-free MCP server declares five tools. Keep the tool schemas, handler routing, + JSON-RPC validation, request IDs, and normalized optional-dependency errors in one module. +- `make lint` is currently blocked by 544 pre-existing Ruff errors under `core/`. `make test` + reaches a missing `pytest` executable and a nonexistent `core/tests/unit` path after the Jest + unit suite passes; use the maintained root Python tests plus npm/Jest gates for this slice. diff --git a/package.json b/package.json index 9ab2d83..fb92579 100644 --- a/package.json +++ b/package.json @@ -70,7 +70,14 @@ "onCommand:openqc.validate", "onView:openqc-sidebar", "onCommand:openqc.previewInput", - "onCommand:openqc.convertFormat" + "onCommand:openqc.convertFormat", + "onCommand:openqc.aiOptimizeInput", + "onCommand:openqc.aiGenerateInput", + "onCommand:openqc.aiExplainParameters", + "onCommand:openqc.aiDebugCalculation", + "onCommand:openqc.aiSettings", + "onCommand:openqc.aiSetApiKey", + "onCommand:openqc.aiClearApiKey" ], "main": "./out/extension.js", "extensionKind": [ @@ -608,6 +615,18 @@ "category": "AI", "icon": "$(gear)" }, + { + "command": "openqc.aiSetApiKey", + "title": "OpenQC: Set OpenAI API Key", + "category": "AI", + "icon": "$(key)" + }, + { + "command": "openqc.aiClearApiKey", + "title": "OpenQC: Clear OpenAI API Key", + "category": "AI", + "icon": "$(trash)" + }, { "command": "openqc.showConverterPanel", "title": "Show ASE Converter", @@ -1377,15 +1396,15 @@ "default": "ollama", "description": "AI provider (OpenAI or Ollama)" }, - "openqc.ai.apiKey": { - "type": "string", - "default": "", - "description": "OpenAI API key (only used with OpenAI provider)" - }, "openqc.ai.model": { "type": "string", "default": "llama2", - "description": "AI model name (gpt-4 for OpenAI, llama2 for Ollama)" + "description": "AI model name for the selected provider" + }, + "openqc.ai.openaiBaseUrl": { + "type": "string", + "default": "https://api.openai.com/v1", + "description": "OpenAI-compatible API base URL" }, "openqc.ai.ollamaUrl": { "type": "string", @@ -1395,8 +1414,24 @@ "openqc.ai.maxTokens": { "type": "number", "default": 2048, + "minimum": 1, + "maximum": 131072, "description": "Maximum tokens in AI response" }, + "openqc.ai.maxOutputChars": { + "type": "number", + "default": 100000, + "minimum": 256, + "maximum": 1000000, + "description": "Maximum AI response characters accepted by the extension" + }, + "openqc.ai.timeoutSeconds": { + "type": "number", + "default": 120, + "exclusiveMinimum": 0, + "maximum": 600, + "description": "Timeout for AI provider and bridge requests" + }, "openqc.ai.temperature": { "type": "number", "default": 0.7, diff --git a/python/openqc/ai/__init__.py b/python/openqc/ai/__init__.py index 86f77ec..97b9214 100644 --- a/python/openqc/ai/__init__.py +++ b/python/openqc/ai/__init__.py @@ -5,8 +5,16 @@ generation, explanation, and debugging. """ -from .client import AIClient, AIRequest, AIResponse -from .prompts import PromptTemplates from .parser import ResponseParser +from .prompts import PromptTemplates + +__all__ = ["AIClient", "AIRequest", "AIResponse", "PromptTemplates", "ResponseParser"] + + +def __getattr__(name): + """Load client types lazily so ``python -m openqc.ai.client`` stays warning-free.""" + if name in {"AIClient", "AIRequest", "AIResponse"}: + from . import client -__all__ = ['AIClient', 'AIRequest', 'AIResponse', 'PromptTemplates', 'ResponseParser'] + return getattr(client, name) + raise AttributeError(name) diff --git a/python/openqc/ai/client.py b/python/openqc/ai/client.py index 4526349..a263ce0 100644 --- a/python/openqc/ai/client.py +++ b/python/openqc/ai/client.py @@ -1,22 +1,29 @@ -""" -AI Client for LLM Integration +"""Dependency-free OpenAI Responses API and Ollama client for OpenQC.""" -Supports OpenAI API and Ollama (local models). -""" +from __future__ import annotations +import argparse +import json import os +import re +import socket import sys -import json -import argparse -from typing import Dict, Optional, Any -from dataclasses import dataclass -from .prompts import PromptTemplates +import urllib.error +import urllib.request +from dataclasses import asdict, dataclass +from typing import Any, Dict, Optional + from .parser import ResponseParser +from .prompts import PromptTemplates + +DEFAULT_TIMEOUT_SECONDS = 120.0 +DEFAULT_MAX_OUTPUT_CHARS = 100_000 @dataclass class AIRequest: """AI request data structure.""" + type: str content: str software: Optional[str] = None @@ -25,7 +32,8 @@ class AIRequest: @dataclass class AIResponse: - """AI response data structure.""" + """AI response returned to the extension process.""" + success: bool content: Optional[str] = None suggestions: Optional[list] = None @@ -34,274 +42,302 @@ class AIResponse: metadata: Optional[Dict[str, Any]] = None +class ProviderError(RuntimeError): + """A provider response that is safe to surface after sanitization.""" + + class AIClient: - """Client for AI/LLM operations.""" - - def __init__(self): - self.provider = os.environ.get('OPENQC_AI_PROVIDER', 'ollama') - self.model = os.environ.get('OPENQC_AI_MODEL', 'llama2') - self.api_key = os.environ.get('OPENQC_AI_API_KEY', '') - self.ollama_url = os.environ.get('OPENQC_AI_OLLAMA_URL', 'http://localhost:11434') - self.max_tokens = int(os.environ.get('OPENQC_AI_MAX_TOKENS', '2048')) - self.temperature = float(os.environ.get('OPENQC_AI_TEMPERATURE', '0.7')) + """Client for OpenAI Responses API and Ollama operations.""" + + def __init__(self) -> None: + self.provider = os.environ.get("OPENQC_AI_PROVIDER", "ollama").strip().lower() + self.model = os.environ.get("OPENQC_AI_MODEL", "llama2").strip() + self.api_key = os.environ.get("OPENQC_AI_API_KEY", "") + self.openai_url = os.environ.get( + "OPENQC_AI_OPENAI_URL", "https://api.openai.com/v1" + ).rstrip("/") + self.ollama_url = os.environ.get( + "OPENQC_AI_OLLAMA_URL", "http://localhost:11434" + ).rstrip("/") + self.max_tokens = _bounded_int("OPENQC_AI_MAX_TOKENS", 2048, 1, 131_072) + self.max_output_chars = _bounded_int( + "OPENQC_AI_MAX_OUTPUT_CHARS", DEFAULT_MAX_OUTPUT_CHARS, 256, 1_000_000 + ) + self.temperature = _bounded_float("OPENQC_AI_TEMPERATURE", 0.7, 0.0, 2.0) + self.timeout = _bounded_float( + "OPENQC_AI_TIMEOUT_SECONDS", DEFAULT_TIMEOUT_SECONDS, 0.1, 600.0 + ) self.parser = ResponseParser() def check(self) -> AIResponse: - """Check if AI backend is available.""" - try: - if self.provider == 'openai': - return self._check_openai() - elif self.provider == 'ollama': - return self._check_ollama() - else: - return AIResponse(success=False, error=f"Unknown provider: {self.provider}") - except Exception as e: - return AIResponse(success=False, error=str(e)) - - def _check_openai(self) -> AIResponse: - """Check OpenAI availability.""" - if not self.api_key: - return AIResponse(success=False, error="OpenAI API key not configured") - - try: - import openai - openai.api_key = self.api_key - # Test with a simple request - models = openai.Model.list() - return AIResponse( - success=True, - content=f"OpenAI connection OK. Available models: {len(models.data)}" - ) - except ImportError: - return AIResponse(success=False, error="OpenAI package not installed. Run: pip install openai") - except Exception as e: - return AIResponse(success=False, error=f"OpenAI connection failed: {e}") - - def _check_ollama(self) -> AIResponse: - """Check Ollama availability.""" - try: - import requests - response = requests.get(f"{self.ollama_url}/api/tags", timeout=5) - if response.status_code == 200: - data = response.json() - models = data.get('models', []) + """Check whether the configured provider is reachable.""" + if self.provider == "openai": + if not self.api_key: return AIResponse( - success=True, - content=f"Ollama connection OK. Available models: {len(models)}" + success=False, error="OpenAI API key is not configured" ) - else: - return AIResponse(success=False, error=f"Ollama returned status {response.status_code}") - except ImportError: - return AIResponse(success=False, error="requests package not installed") - except Exception as e: - return AIResponse(success=False, error=f"Ollama connection failed: {e}") + return self._request_status( + f"{self.openai_url}/models", + headers={"Authorization": f"Bearer {self.api_key}"}, + provider="OpenAI", + ) + if self.provider == "ollama": + return self._request_status( + f"{self.ollama_url}/api/tags", provider="Ollama" + ) + return AIResponse(success=False, error=f"Unknown AI provider: {self.provider}") def optimize(self, request: AIRequest) -> AIResponse: - """Optimize input file.""" + """Optimize an input file.""" if not request.software: return AIResponse(success=False, error="Software type not specified") - - prompt = PromptTemplates.get_optimize_prompt( - request.content, - request.software, - request.context + response = self._call_llm( + PromptTemplates.get_optimize_prompt( + request.content, request.software, request.context + ) ) - - response = self._call_llm(prompt) if not response.success: return response - - # Parse suggestions from response - parsed = self.parser.parse_optimization(response.content or '') + parsed = self.parser.parse_optimization(response.content or "") return AIResponse( success=True, - content=parsed.get('content'), - suggestions=parsed.get('suggestions', []), - metadata={'raw_response': response.content} + content=parsed.get("content"), + suggestions=parsed.get("suggestions", []), + metadata=response.metadata, ) def generate(self, request: AIRequest) -> AIResponse: - """Generate input file.""" + """Generate an input file.""" if not request.software: return AIResponse(success=False, error="Software type not specified") - - prompt = PromptTemplates.get_generate_prompt( - request.content, - request.software, - request.context + response = self._call_llm( + PromptTemplates.get_generate_prompt( + request.content, request.software, request.context + ) ) - - response = self._call_llm(prompt) if not response.success: return response - - # Extract input file from response - generated = self.parser.parse_generated_input(response.content or '') return AIResponse( success=True, - generatedInput=generated, - metadata={'raw_response': response.content} + generatedInput=self.parser.parse_generated_input(response.content or ""), + metadata=response.metadata, ) def explain(self, request: AIRequest) -> AIResponse: - """Explain parameters.""" + """Explain input parameters.""" if not request.software: return AIResponse(success=False, error="Software type not specified") - - prompt = PromptTemplates.get_explain_prompt( - request.content, - request.software + return self._call_llm( + PromptTemplates.get_explain_prompt(request.content, request.software) ) - - return self._call_llm(prompt) def debug(self, request: AIRequest) -> AIResponse: - """Debug calculation.""" + """Debug a failed calculation.""" if not request.software: return AIResponse(success=False, error="Software type not specified") - - output = request.context.get('output', '') if request.context else '' - - prompt = PromptTemplates.get_debug_prompt( - request.content, - output, - request.software + output = request.context.get("output", "") if request.context else "" + return self._call_llm( + PromptTemplates.get_debug_prompt(request.content, output, request.software) ) - - return self._call_llm(prompt) - def _call_llm(self, prompt: str) -> AIResponse: - """Call the LLM with the given prompt.""" + def _request_status( + self, url: str, provider: str, headers: Optional[Dict[str, str]] = None + ) -> AIResponse: try: - if self.provider == 'openai': - return self._call_openai(prompt) - elif self.provider == 'ollama': - return self._call_ollama(prompt) - else: - return AIResponse(success=False, error=f"Unknown provider: {self.provider}") - except Exception as e: - return AIResponse(success=False, error=f"LLM call failed: {e}") + self._http_json(url, headers=headers) + return AIResponse(success=True, content=f"{provider} connection OK") + except ( + Exception + ) as exc: # Boundary converts all network failures to protocol JSON. + return AIResponse(success=False, error=self._safe_error(provider, exc)) + + def _call_llm(self, prompt: str) -> AIResponse: + if self.provider == "openai": + return self._call_openai(prompt) + if self.provider == "ollama": + return self._call_ollama(prompt) + return AIResponse(success=False, error=f"Unknown AI provider: {self.provider}") def _call_openai(self, prompt: str) -> AIResponse: - """Call OpenAI API.""" + if not self.api_key: + return AIResponse(success=False, error="OpenAI API key is not configured") + payload = { + "model": self.model, + "instructions": "You are a quantum chemistry expert.", + "input": prompt, + "max_output_tokens": self.max_tokens, + "temperature": self.temperature, + } try: - import openai - openai.api_key = self.api_key - - response = openai.ChatCompletion.create( - model=self.model, - messages=[ - {"role": "system", "content": "You are a quantum chemistry expert."}, - {"role": "user", "content": prompt} - ], - max_tokens=self.max_tokens, - temperature=self.temperature - ) - - content = response.choices[0].message.content - return AIResponse(success=True, content=content) - except ImportError: - return AIResponse( - success=False, - error="OpenAI package not installed. Run: pip install openai" + data = self._http_json( + f"{self.openai_url}/responses", + payload, + {"Authorization": f"Bearer {self.api_key}"}, ) - except Exception as e: - return AIResponse(success=False, error=f"OpenAI API error: {e}") + content = _extract_openai_text(data) + return self._bounded_success(content) + except ( + Exception + ) as exc: # Boundary converts all provider failures to protocol JSON. + return AIResponse(success=False, error=self._safe_error("OpenAI", exc)) def _call_ollama(self, prompt: str) -> AIResponse: - """Call Ollama API.""" + payload = { + "model": self.model, + "prompt": prompt, + "stream": False, + "options": { + "temperature": self.temperature, + "num_predict": self.max_tokens, + }, + } try: - import requests - - response = requests.post( - f"{self.ollama_url}/api/generate", - json={ - "model": self.model, - "prompt": prompt, - "stream": False, - "options": { - "temperature": self.temperature, - "num_predict": self.max_tokens - } - }, - timeout=120 - ) - - if response.status_code == 200: - data = response.json() - content = data.get('response', '') - return AIResponse(success=True, content=content) - else: - return AIResponse( - success=False, - error=f"Ollama returned status {response.status_code}: {response.text}" - ) - except ImportError: - return AIResponse( - success=False, - error="requests package not installed. Run: pip install requests" - ) - except Exception as e: - return AIResponse(success=False, error=f"Ollama API error: {e}") + data = self._http_json(f"{self.ollama_url}/api/generate", payload) + content = data.get("response") + if not isinstance(content, str): + raise ProviderError("provider response did not contain text") + return self._bounded_success(content) + except ( + Exception + ) as exc: # Boundary converts all provider failures to protocol JSON. + return AIResponse(success=False, error=self._safe_error("Ollama", exc)) + + def _bounded_success(self, content: str) -> AIResponse: + truncated = len(content) > self.max_output_chars + bounded = content[: self.max_output_chars] + metadata = {"truncated": True} if truncated else None + return AIResponse(success=True, content=bounded, metadata=metadata) + + def _http_json( + self, + url: str, + payload: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, str]] = None, + ) -> Dict[str, Any]: + request_headers = {"Accept": "application/json", "User-Agent": "OpenQC/0.0.1"} + request_headers.update(headers or {}) + body = None + if payload is not None: + body = json.dumps(payload).encode("utf-8") + request_headers["Content-Type"] = "application/json" + request = urllib.request.Request(url, data=body, headers=request_headers) + try: + with urllib.request.urlopen(request, timeout=self.timeout) as response: + raw = response.read(self.max_output_chars * 4 + 1) + if len(raw) > self.max_output_chars * 4: + raise ProviderError( + "provider response exceeded the safe size limit" + ) + result = json.loads(raw.decode("utf-8")) + if not isinstance(result, dict): + raise ProviderError("provider returned invalid JSON") + return result + except urllib.error.HTTPError as exc: + detail = "" + try: + body_data = json.loads(exc.read(4096).decode("utf-8", errors="replace")) + detail = _provider_error_message(body_data) + except (json.JSONDecodeError, UnicodeDecodeError): + pass + finally: + exc.close() + suffix = f": {detail}" if detail else "" + raise ProviderError(f"HTTP {exc.code}{suffix}") from None + except (urllib.error.URLError, socket.timeout, TimeoutError) as exc: + reason = getattr(exc, "reason", exc) + if ( + isinstance(reason, (socket.timeout, TimeoutError)) + or "timed out" in str(reason).lower() + ): + raise ProviderError( + f"request timed out after {self.timeout:g} seconds" + ) from None + raise ProviderError("provider is unavailable") from None + except json.JSONDecodeError: + raise ProviderError("provider returned invalid JSON") from None + + def _safe_error(self, provider: str, error: Exception) -> str: + message = str(error) or error.__class__.__name__ + if self.api_key: + message = message.replace(self.api_key, "[REDACTED]") + message = re.sub(r"(?i)bearer\s+[^\s,;]+", "Bearer [REDACTED]", message) + message = re.sub(r"\bsk-[A-Za-z0-9_-]{8,}\b", "[REDACTED]", message) + return f"{provider} request failed: {message[:500]}" + + +def _extract_openai_text(data: Dict[str, Any]) -> str: + output_text = data.get("output_text") + if isinstance(output_text, str): + return output_text + parts = [] + for item in data.get("output", []): + if not isinstance(item, dict): + continue + for content in item.get("content", []): + if isinstance(content, dict) and isinstance(content.get("text"), str): + parts.append(content["text"]) + if not parts: + raise ProviderError("provider response did not contain output text") + return "".join(parts) + + +def _provider_error_message(data: Any) -> str: + if not isinstance(data, dict): + return "" + error = data.get("error", data) + if isinstance(error, dict) and isinstance(error.get("message"), str): + return error["message"][:300] + if isinstance(error, str): + return error[:300] + return "" -def main(): - """CLI entry point.""" - parser = argparse.ArgumentParser(description='OpenQC AI Client') - parser.add_argument('command', choices=['check', 'optimize', 'generate', 'explain', 'debug']) +def _bounded_int(name: str, default: int, minimum: int, maximum: int) -> int: + try: + value = int(os.environ.get(name, str(default))) + except ValueError: + value = default + return max(minimum, min(maximum, value)) + + +def _bounded_float(name: str, default: float, minimum: float, maximum: float) -> float: + try: + value = float(os.environ.get(name, str(default))) + except ValueError: + value = default + return max(minimum, min(maximum, value)) + + +def main() -> None: + """CLI/module entry point.""" + parser = argparse.ArgumentParser(description="OpenQC AI Client") + parser.add_argument( + "command", choices=["check", "optimize", "generate", "explain", "debug"] + ) args = parser.parse_args() - client = AIClient() - - if args.command == 'check': - response = client.check() - print(json.dumps({ - 'success': response.success, - 'content': response.content, - 'error': response.error - })) - return - - # Read request from stdin - request_data = sys.stdin.read() - if not request_data: - print(json.dumps({'success': False, 'error': 'No input provided'})) + + if args.command == "check": + print(json.dumps(asdict(client.check()))) return - + try: - data = json.loads(request_data) + data = json.loads(sys.stdin.read()) + if not isinstance(data, dict): + raise ValueError("request must be a JSON object") request = AIRequest( - type=data.get('type', ''), - content=data.get('content', ''), - software=data.get('software'), - context=data.get('context') + type=str(data.get("type", "")), + content=str(data.get("content", "")), + software=data.get("software"), + context=data.get("context"), ) - except json.JSONDecodeError as e: - print(json.dumps({'success': False, 'error': f'Invalid JSON: {e}'})) - return - - # Execute command - if args.command == 'optimize': - response = client.optimize(request) - elif args.command == 'generate': - response = client.generate(request) - elif args.command == 'explain': - response = client.explain(request) - elif args.command == 'debug': - response = client.debug(request) - else: - response = AIResponse(success=False, error=f"Unknown command: {args.command}") - - # Output response - print(json.dumps({ - 'success': response.success, - 'content': response.content, - 'suggestions': response.suggestions, - 'generatedInput': response.generatedInput, - 'error': response.error, - 'metadata': response.metadata - })) - - -if __name__ == '__main__': + operation = getattr(client, args.command) + response = operation(request) + except (json.JSONDecodeError, TypeError, ValueError) as exc: + response = AIResponse(success=False, error=f"Invalid request: {str(exc)[:300]}") + + print(json.dumps(asdict(response))) + + +if __name__ == "__main__": main() diff --git a/python/openqc/mcp_server.py b/python/openqc/mcp_server.py index 6e1b25f..7c3eb90 100644 --- a/python/openqc/mcp_server.py +++ b/python/openqc/mcp_server.py @@ -1,66 +1,59 @@ -""" -OpenQC MCP Server - Exposes OpenQC operations as MCP tools. +"""Dependency-free JSON-RPC 2.0 MCP server for OpenQC scientific bridges.""" -This is a minimal MCP server that lists available tools and can be -launched by MCP-compatible clients (Claude, Codex, Cursor, etc.). - -Usage: - python -m openqc.mcp_server - -Protocol: JSON-RPC 2.0 over stdin/stdout (MCP standard transport) -""" +from __future__ import annotations import json +import re import sys - +from typing import Any, Callable, Dict, Optional TOOLS = [ { "name": "openqc.parseStructure", - "description": "Parse a molecular/crystal structure file into OpenQCStructure JSON", + "description": "Parse a molecular or crystal structure into OpenQCStructure JSON", "inputSchema": { "type": "object", "properties": { "path": {"type": "string", "description": "File path"}, - "format": {"type": "string", "description": "Format hint (auto, xyz, poscar, cif)"}, + "format": {"type": "string", "description": "Optional format hint"}, }, "required": ["path"], }, }, { "name": "openqc.parseCalculationOutput", - "description": "Parse a quantum chemistry output file using cclib", + "description": "Parse calculation output with cclib and a native fallback", "inputSchema": { "type": "object", "properties": { "path": {"type": "string", "description": "Output file path"}, - "software": {"type": "string", "description": "Software hint"}, + "software": {"type": "string", "description": "Optional software hint"}, }, "required": ["path"], }, }, { "name": "openqc.checkPythonBackend", - "description": "Check Scientific Python backend status", + "description": "Check the Scientific Python backend and optional dependencies", "inputSchema": {"type": "object", "properties": {}}, }, { "name": "openqc.generateSupercell", - "description": "Generate a supercell preview", + "description": "Generate a bounded supercell preview", "inputSchema": { "type": "object", "properties": { "path": {"type": "string"}, - "na": {"type": "number", "default": 2}, - "nb": {"type": "number", "default": 2}, - "nc": {"type": "number", "default": 2}, + "na": {"type": "integer", "minimum": 1, "maximum": 100, "default": 2}, + "nb": {"type": "integer", "minimum": 1, "maximum": 100, "default": 2}, + "nc": {"type": "integer", "minimum": 1, "maximum": 100, "default": 2}, }, "required": ["path"], }, }, { "name": "openqc.summarizeDataset", - "description": "Summarize an atomistic dataset (dpdata)", + "description": "Summarize an atomistic dataset when dpdata is installed", "inputSchema": { "type": "object", "properties": { @@ -72,96 +65,230 @@ }, ] +_TOOL_SCHEMAS = {tool["name"]: tool["inputSchema"] for tool in TOOLS} + + +class RpcError(Exception): + """JSON-RPC error with a stable code and safe message.""" + + def __init__(self, code: int, message: str): + super().__init__(message) + self.code = code + self.message = message + + +def _parse_structure(arguments: Dict[str, Any]) -> Dict[str, Any]: + from openqc.bridge.structure_bridge import parse_file + + return parse_file(arguments["path"], arguments.get("format")) + + +def _parse_output(arguments: Dict[str, Any]) -> Dict[str, Any]: + from openqc.bridge.output_bridge import _parse_output_native, _parse_with_cclib + + path = arguments["path"] + software = arguments.get("software", "auto") + try: + result = _parse_with_cclib(path, software) + if result is None: + raise ValueError("cclib returned no output") + result["cclibAvailable"] = True + return result + except Exception as cclib_error: + with open(path, "r", encoding="utf-8", errors="ignore") as handle: + result = _parse_output_native(handle.read(), software) + result["sourceFile"] = path + result["cclibAvailable"] = False + result.setdefault("warnings", []).append( + f"cclib parser unavailable or failed: {_safe_message(cclib_error)}" + ) + return result + + +def _check_backend(_arguments: Dict[str, Any]) -> Dict[str, Any]: + from openqc.bridge.check_backend import check_backend + + return check_backend() + + +def _generate_supercell(arguments: Dict[str, Any]) -> Dict[str, Any]: + from openqc.bridge.structure_bridge import _generate_supercell, parse_file + + structure = parse_file(arguments["path"]) + return _generate_supercell( + structure, + arguments.get("na", 2), + arguments.get("nb", 2), + arguments.get("nc", 2), + ) + + +def _summarize_dataset(arguments: Dict[str, Any]) -> Dict[str, Any]: + from openqc.bridge.dpdata_bridge import _summarize_dataset + + return _summarize_dataset(arguments["path"], arguments.get("format", "auto")) + + +_HANDLERS: Dict[str, Callable[[Dict[str, Any]], Dict[str, Any]]] = { + "openqc.parseStructure": _parse_structure, + "openqc.parseCalculationOutput": _parse_output, + "openqc.checkPythonBackend": _check_backend, + "openqc.generateSupercell": _generate_supercell, + "openqc.summarizeDataset": _summarize_dataset, +} -def handle_request(request: dict) -> dict: - """Handle a JSON-RPC request.""" - method = request.get("method", "") - req_id = request.get("id") + +def _validate_request(request: Any) -> Dict[str, Any]: + if not isinstance(request, dict) or request.get("jsonrpc") != "2.0": + raise RpcError(-32600, "Invalid Request") + if not isinstance(request.get("method"), str): + raise RpcError(-32600, "Invalid Request") params = request.get("params", {}) + if not isinstance(params, dict): + raise RpcError(-32602, "Invalid params: params must be an object") + return request - if method == "initialize": - return { - "jsonrpc": "2.0", - "id": req_id, - "result": { - "protocolVersion": "2024-11-05", - "capabilities": {"tools": {}}, - "serverInfo": {"name": "openqc-mcp", "version": "0.1.0"}, - }, - } - elif method == "tools/list": - return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}} +def _validate_arguments(tool_name: str, arguments: Any) -> Dict[str, Any]: + if not isinstance(arguments, dict): + raise RpcError(-32602, "Invalid params: arguments must be an object") + schema = _TOOL_SCHEMAS[tool_name] + result = dict(arguments) + for name in schema.get("required", []): + if name not in result: + raise RpcError( + -32602, f"Invalid params: missing required argument '{name}'" + ) + for name, definition in schema.get("properties", {}).items(): + if name not in result and "default" in definition: + result[name] = definition["default"] + if name not in result: + continue + value = result[name] + expected = definition.get("type") + if expected == "string" and (not isinstance(value, str) or not value.strip()): + raise RpcError( + -32602, f"Invalid params: '{name}' must be a non-empty string" + ) + if expected == "integer" and ( + not isinstance(value, int) or isinstance(value, bool) + ): + raise RpcError(-32602, f"Invalid params: '{name}' must be an integer") + minimum = definition.get("minimum") + maximum = definition.get("maximum") + if minimum is not None and value < minimum: + raise RpcError( + -32602, f"Invalid params: '{name}' is below minimum {minimum}" + ) + if maximum is not None and value > maximum: + raise RpcError( + -32602, f"Invalid params: '{name}' exceeds maximum {maximum}" + ) + return result - elif method == "tools/call": - tool_name = params.get("name", "") - arguments = params.get("arguments", {}) - # Route to appropriate bridge - if tool_name == "openqc.checkPythonBackend": - from openqc.bridge.check_backend import check_backend - result = check_backend() - return {"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": json.dumps(result)}]}} +def _result(request_id: Any, payload: Dict[str, Any]) -> Dict[str, Any]: + return { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "content": [ + {"type": "text", "text": json.dumps(payload, ensure_ascii=False)} + ] + }, + } - elif tool_name == "openqc.parseStructure": - from openqc.bridge.structure_bridge import parse_file - result = parse_file(arguments["path"], arguments.get("format")) - return {"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": json.dumps(result)}]}} - elif tool_name == "openqc.parseCalculationOutput": - from openqc.bridge.output_bridge import _parse_with_cclib - try: - result = _parse_with_cclib(arguments["path"], arguments.get("software")) - return {"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": json.dumps(result)}]}} - except ImportError as e: - return {"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": json.dumps({"error": str(e)})}]}} +def _error(request_id: Any, code: int, message: str) -> Dict[str, Any]: + return { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": code, "message": message}, + } - else: - return { + +def _safe_message(error: Exception) -> str: + message = str(error) or error.__class__.__name__ + message = re.sub(r"(?i)bearer\s+[^\s,;]+", "Bearer [REDACTED]", message) + message = re.sub(r"\bsk-[A-Za-z0-9_-]{8,}\b", "[REDACTED]", message) + return message[:500] + + +def handle_request(request: Any) -> Optional[Dict[str, Any]]: + """Validate and dispatch one JSON-RPC request without terminating the server.""" + request_id = request.get("id") if isinstance(request, dict) else None + is_notification = isinstance(request, dict) and "id" not in request + try: + validated = _validate_request(request) + method = validated["method"] + params = validated.get("params", {}) + + if method == "notifications/initialized": + return None + if method == "initialize": + response = { "jsonrpc": "2.0", - "id": req_id, - "error": {"code": -32601, "message": f"Unknown tool: {tool_name}"}, + "id": request_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "openqc-mcp", "version": "0.1.0"}, + }, } + elif method == "tools/list": + response = { + "jsonrpc": "2.0", + "id": request_id, + "result": {"tools": TOOLS}, + } + elif method == "tools/call": + tool_name = params.get("name") + if not isinstance(tool_name, str): + raise RpcError(-32602, "Invalid params: tool name must be a string") + handler = _HANDLERS.get(tool_name) + if handler is None: + raise RpcError(-32601, f"Unknown tool: {tool_name}") + arguments = _validate_arguments(tool_name, params.get("arguments", {})) + try: + response = _result(request_id, handler(arguments)) + except ImportError as exc: + raise RpcError( + -32001, f"Optional dependency unavailable: {_safe_message(exc)}" + ) from None + except ( + FileNotFoundError, + IsADirectoryError, + PermissionError, + ValueError, + ) as exc: + raise RpcError( + -32002, f"Tool execution failed: {_safe_message(exc)}" + ) from None + except Exception as exc: + raise RpcError( + -32603, f"Internal error: {_safe_message(exc)}" + ) from None + else: + raise RpcError(-32601, f"Unknown method: {method}") - elif method == "notifications/initialized": - return None # No response for notifications - - else: - return { - "jsonrpc": "2.0", - "id": req_id, - "error": {"code": -32601, "message": f"Unknown method: {method}"}, - } + return None if is_notification else response + except RpcError as exc: + return None if is_notification else _error(request_id, exc.code, exc.message) -def main(): - """Run the MCP server.""" +def main() -> None: + """Run newline-delimited JSON-RPC over stdin/stdout.""" for line in sys.stdin: - line = line.strip() - if not line: + if not line.strip(): continue - try: request = json.loads(line) + except json.JSONDecodeError: + response = _error(None, -32700, "Parse error") + else: response = handle_request(request) - if response is not None: - sys.stdout.write(json.dumps(response) + "\n") - sys.stdout.flush() - except json.JSONDecodeError as e: - error_response = { - "jsonrpc": "2.0", - "id": None, - "error": {"code": -32700, "message": f"Parse error: {e}"}, - } - sys.stdout.write(json.dumps(error_response) + "\n") - sys.stdout.flush() - except Exception as e: - error_response = { - "jsonrpc": "2.0", - "id": None, - "error": {"code": -32603, "message": f"Internal error: {e}"}, - } - sys.stdout.write(json.dumps(error_response) + "\n") + if response is not None: + sys.stdout.write(json.dumps(response, ensure_ascii=False) + "\n") sys.stdout.flush() diff --git a/src/ai/AICore.ts b/src/ai/AICore.ts index dfdf96d..7db7b3a 100644 --- a/src/ai/AICore.ts +++ b/src/ai/AICore.ts @@ -1,38 +1,29 @@ -/** - * AI Core Module - TypeScript Interface - * - * Provides TypeScript interface to the Python AI backend for - * AI-powered input file optimization, generation, and debugging. - */ +/** Secure TypeScript bridge to the Python AI module. */ -import * as vscode from 'vscode'; import * as path from 'path'; import { spawn } from 'child_process'; +import * as vscode from 'vscode'; + +export const OPENAI_API_KEY_SECRET = 'openqc.ai.openaiApiKey'; +const STDERR_LIMIT = 8_192; -/** - * Supported AI providers - */ export enum AIProvider { OpenAI = 'openai', Ollama = 'ollama', } -/** - * AI Configuration interface - */ export interface AIConfig { enabled: boolean; provider: AIProvider; - apiKey?: string; model: string; + openaiBaseUrl: string; ollamaUrl: string; maxTokens: number; + maxOutputChars: number; + timeoutSeconds: number; temperature: number; } -/** - * AI Request types - */ export enum AIRequestType { OptimizeInput = 'optimize_input', GenerateInput = 'generate_input', @@ -40,19 +31,13 @@ export enum AIRequestType { DebugCalculation = 'debug_calculation', } -/** - * AI Request interface - */ export interface AIRequest { type: AIRequestType; content: string; software?: string; - context?: Record; + context?: Record; } -/** - * AI Response interface - */ export interface AIResponse { success: boolean; content?: string; @@ -66,266 +51,321 @@ export interface AIResponse { }>; generatedInput?: string; error?: string; - metadata?: Record; + metadata?: Record; } -/** - * AI Core class - * - * Interfaces with Python AI backend for LLM operations - */ export class AICore { private config: AIConfig; private pythonPath: string; - private aiScript: string; - private context: vscode.ExtensionContext; + private readonly pythonRoot: string; - constructor(context: vscode.ExtensionContext) { - this.context = context; + constructor(private readonly context: vscode.ExtensionContext) { this.pythonPath = this.getPythonPath(); - this.aiScript = path.join(context.extensionPath, 'python', 'openqc', 'ai', 'client.py'); + this.pythonRoot = path.join(context.extensionPath, 'python'); this.config = this.loadConfig(); } - /** - * Get Python executable path from configuration - */ private getPythonPath(): string { - const config = vscode.workspace.getConfiguration('openqc'); - return config.get('pythonPath', 'python3'); + return vscode.workspace.getConfiguration('openqc').get('pythonPath', 'python3'); } - /** - * Load AI configuration from VSCode settings - */ private loadConfig(): AIConfig { const config = vscode.workspace.getConfiguration('openqc.ai'); return { enabled: config.get('enabled', false), provider: config.get('provider', AIProvider.Ollama), - apiKey: config.get('apiKey'), model: config.get('model', 'llama2'), + openaiBaseUrl: config.get('openaiBaseUrl', 'https://api.openai.com/v1'), ollamaUrl: config.get('ollamaUrl', 'http://localhost:11434'), maxTokens: config.get('maxTokens', 2048), + maxOutputChars: config.get('maxOutputChars', 100_000), + timeoutSeconds: config.get('timeoutSeconds', 120), temperature: config.get('temperature', 0.7), }; } - /** - * Refresh configuration (call when settings change) - */ public refreshConfig(): void { + this.pythonPath = this.getPythonPath(); this.config = this.loadConfig(); } - /** - * Check if AI features are enabled and available - */ public isEnabled(): boolean { return this.config.enabled; } - /** - * Check if AI backend is available - */ - public async isAvailable(): Promise { - if (!this.config.enabled) { - return false; + public async setOpenAIApiKey(apiKey: string): Promise { + const value = apiKey.trim(); + if (!value) { + throw new Error('OpenAI API key cannot be empty'); } + await this.context.secrets.store(OPENAI_API_KEY_SECRET, value); + } - try { - const result = await this.executeAI(['check'], ''); - return result.success; - } catch { + public async clearOpenAIApiKey(): Promise { + await this.context.secrets.delete(OPENAI_API_KEY_SECRET); + } + + public async hasOpenAIApiKey(): Promise { + return Boolean(await this.context.secrets.get(OPENAI_API_KEY_SECRET)); + } + + public async isAvailable(token?: vscode.CancellationToken): Promise { + if (!this.config.enabled) { return false; } + const result = await this.executeAI(['check'], '', token); + return result.success; } - /** - * Get current configuration - */ public getConfig(): AIConfig { return { ...this.config }; } - /** - * Optimize input file parameters - */ - public async optimizeInput(inputContent: string, software: string): Promise { - const request: AIRequest = { - type: AIRequestType.OptimizeInput, - content: inputContent, - software, - }; - - return this.executeAI(['optimize'], JSON.stringify(request)); + public optimizeInput( + inputContent: string, + software: string, + token?: vscode.CancellationToken + ): Promise { + return this.executeAI( + ['optimize'], + JSON.stringify({ + type: AIRequestType.OptimizeInput, + content: inputContent, + software, + } satisfies AIRequest), + token + ); } - /** - * Generate input file from natural language description - */ - public async generateInput( + public generateInput( description: string, software: string, - context?: Record + context?: Record, + token?: vscode.CancellationToken ): Promise { - const request: AIRequest = { - type: AIRequestType.GenerateInput, - content: description, - software, - context, - }; - - return this.executeAI(['generate'], JSON.stringify(request)); + return this.executeAI( + ['generate'], + JSON.stringify({ + type: AIRequestType.GenerateInput, + content: description, + software, + context, + } satisfies AIRequest), + token + ); } - /** - * Explain input file parameters - */ - public async explainParameters(inputContent: string, software: string): Promise { - const request: AIRequest = { - type: AIRequestType.ExplainParameters, - content: inputContent, - software, - }; - - return this.executeAI(['explain'], JSON.stringify(request)); + public explainParameters( + inputContent: string, + software: string, + token?: vscode.CancellationToken + ): Promise { + return this.executeAI( + ['explain'], + JSON.stringify({ + type: AIRequestType.ExplainParameters, + content: inputContent, + software, + } satisfies AIRequest), + token + ); } - /** - * Debug failed calculation - */ - public async debugCalculation( + public debugCalculation( inputContent: string, outputContent: string, - software: string + software: string, + token?: vscode.CancellationToken ): Promise { - const request: AIRequest = { - type: AIRequestType.DebugCalculation, - content: inputContent, - software, - context: { output: outputContent }, - }; - - return this.executeAI(['debug'], JSON.stringify(request)); + return this.executeAI( + ['debug'], + JSON.stringify({ + type: AIRequestType.DebugCalculation, + content: inputContent, + software, + context: { output: outputContent }, + } satisfies AIRequest), + token + ); } - /** - * Execute AI Python script - */ - private async executeAI(args: string[], input: string): Promise { - return new Promise((resolve, reject) => { - const env = { - ...process.env, - OPENQC_AI_PROVIDER: this.config.provider, - OPENQC_AI_MODEL: this.config.model, - OPENQC_AI_API_KEY: this.config.apiKey || '', - OPENQC_AI_OLLAMA_URL: this.config.ollamaUrl, - OPENQC_AI_MAX_TOKENS: this.config.maxTokens.toString(), - OPENQC_AI_TEMPERATURE: this.config.temperature.toString(), - }; + private async executeAI( + args: string[], + input: string, + token?: vscode.CancellationToken + ): Promise { + const apiKey = + this.config.provider === AIProvider.OpenAI + ? await this.context.secrets.get(OPENAI_API_KEY_SECRET) + : undefined; + if (this.config.provider === AIProvider.OpenAI && !apiKey) { + return { success: false, error: 'OpenAI API key is not configured' }; + } + if (token?.isCancellationRequested) { + return { success: false, error: 'AI request was cancelled' }; + } - const process_ = spawn(this.pythonPath, [this.aiScript, ...args], { env }); + return new Promise(resolve => { + const pathSeparator = process.platform === 'win32' ? ';' : ':'; + const currentPythonPath = process.env.PYTHONPATH; + const child = spawn(this.pythonPath, ['-m', 'openqc.ai.client', ...args], { + cwd: this.pythonRoot, + shell: false, + env: { + ...process.env, + PYTHONPATH: currentPythonPath + ? `${this.pythonRoot}${pathSeparator}${currentPythonPath}` + : this.pythonRoot, + OPENQC_AI_PROVIDER: this.config.provider, + OPENQC_AI_MODEL: this.config.model, + OPENQC_AI_API_KEY: apiKey ?? '', + OPENQC_AI_OPENAI_URL: this.config.openaiBaseUrl, + OPENQC_AI_OLLAMA_URL: this.config.ollamaUrl, + OPENQC_AI_MAX_TOKENS: String(this.config.maxTokens), + OPENQC_AI_MAX_OUTPUT_CHARS: String(this.config.maxOutputChars), + OPENQC_AI_TIMEOUT_SECONDS: String(this.config.timeoutSeconds), + OPENQC_AI_TEMPERATURE: String(this.config.temperature), + }, + }); let stdout = ''; let stderr = ''; + let settled = false; + let exceededOutputLimit = false; + const outputLimit = Math.max(4096, this.config.maxOutputChars * 4); + + const finish = (response: AIResponse): void => { + if (settled) { + return; + } + settled = true; + clearTimeout(timeout); + cancellation.dispose(); + resolve(response); + }; + + const terminate = (message: string): void => { + child.kill(); + finish({ success: false, error: message }); + }; + + const timeout = setTimeout( + () => terminate(`AI request timed out after ${this.config.timeoutSeconds} seconds`), + Math.max(100, this.config.timeoutSeconds * 1000) + ); + const cancellation = token?.onCancellationRequested(() => + terminate('AI request was cancelled') + ) ?? { + dispose: () => undefined, + }; if (input) { - process_.stdin.write(input); - process_.stdin.end(); + child.stdin.write(input); } + child.stdin.end(); - process_.stdout.on('data', data => { + child.stdout.on('data', data => { + if (settled || exceededOutputLimit) { + return; + } stdout += data.toString(); + if (Buffer.byteLength(stdout, 'utf8') > outputLimit) { + exceededOutputLimit = true; + terminate('AI backend output exceeded the configured safe limit'); + } }); - process_.stderr.on('data', data => { - stderr += data.toString(); + child.stderr.on('data', data => { + if (stderr.length < STDERR_LIMIT) { + stderr += data.toString().slice(0, STDERR_LIMIT - stderr.length); + } }); - process_.on('close', code => { + child.on('close', code => { + if (settled) { + return; + } if (code !== 0) { - resolve({ + const detail = this.sanitizeError(stderr, apiKey); + finish({ success: false, - error: `AI backend failed with code ${code}: ${stderr}`, + error: detail + ? `AI backend failed with code ${code}: ${detail}` + : `AI backend failed with code ${code}`, }); return; } - try { - const result = JSON.parse(stdout) as AIResponse; - resolve(result); - } catch (error) { - resolve({ - success: false, - error: `Failed to parse AI response: ${error}`, - metadata: { raw_output: stdout }, - }); + finish(JSON.parse(stdout) as AIResponse); + } catch { + finish({ success: false, error: 'AI backend returned an invalid response' }); } }); - process_.on('error', error => { - reject({ + child.on('error', error => { + finish({ success: false, - error: `Failed to execute AI backend: ${error.message}`, + error: `Failed to execute AI backend: ${this.sanitizeError(error.message, apiKey)}`, }); }); }); } - /** - * Get available models for current provider - */ + private sanitizeError(message: string, apiKey?: string): string { + let safe = message; + if (apiKey) { + safe = safe.split(apiKey).join('[REDACTED]'); + } + return safe + .replace(/bearer\s+[^\s,;]+/gi, 'Bearer [REDACTED]') + .replace(/\bsk-[A-Za-z0-9_-]{8,}\b/g, '[REDACTED]') + .trim() + .slice(0, 500); + } + public async getAvailableModels(): Promise { if (this.config.provider === AIProvider.OpenAI) { - return ['gpt-4', 'gpt-4-turbo', 'gpt-3.5-turbo']; - } else if (this.config.provider === AIProvider.Ollama) { - try { - const response = await fetch(`${this.config.ollamaUrl}/api/tags`); - if (!response.ok) { - return ['llama2', 'codellama', 'mistral']; - } - const data = await response.json(); - return data.models?.map((m: any) => m.name) || ['llama2']; - } catch { - return ['llama2', 'codellama', 'mistral']; + return this.config.model ? [this.config.model] : []; + } + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), this.config.timeoutSeconds * 1000); + try { + const response = await fetch(`${this.config.ollamaUrl}/api/tags`, { + signal: controller.signal, + }); + if (!response.ok) { + return []; } + const data = (await response.json()) as { models?: Array<{ name?: string }> }; + return data.models?.flatMap(model => (model.name ? [model.name] : [])) ?? []; + } catch { + return []; + } finally { + clearTimeout(timeout); } - return []; } - /** - * Validate AI configuration - */ public validateConfig(): { valid: boolean; errors: string[] } { const errors: string[] = []; - if (!this.config.enabled) { errors.push('AI features are disabled'); } - - if (this.config.provider === AIProvider.OpenAI && !this.config.apiKey) { - errors.push('OpenAI API key is required'); - } - if (!this.config.model) { errors.push('Model name is required'); } - if (this.config.temperature < 0 || this.config.temperature > 2) { errors.push('Temperature must be between 0 and 2'); } - - return { - valid: errors.length === 0, - errors, - }; + if (this.config.timeoutSeconds <= 0 || this.config.timeoutSeconds > 600) { + errors.push('Timeout must be between 0 and 600 seconds'); + } + if (this.config.maxOutputChars < 256 || this.config.maxOutputChars > 1_000_000) { + errors.push('Maximum output characters must be between 256 and 1000000'); + } + return { valid: errors.length === 0, errors }; } } -/** - * AI Core factory for creating instances - */ export class AICoreFactory { private static instance: AICore | null = null; diff --git a/src/ai/aiCommands.ts b/src/ai/aiCommands.ts index 97f7a99..2a769bb 100644 --- a/src/ai/aiCommands.ts +++ b/src/ai/aiCommands.ts @@ -7,7 +7,7 @@ import * as vscode from 'vscode'; import * as fs from 'fs'; -import { AICore, AICoreFactory, AIResponse } from './AICore'; +import { AICore, AICoreFactory, AIProvider, AIResponse } from './AICore'; /** * Register AI commands @@ -45,6 +45,29 @@ export function registerAICommands(context: vscode.ExtensionContext): void { }) ); + context.subscriptions.push( + vscode.commands.registerCommand('openqc.aiSetApiKey', async () => { + const apiKey = await vscode.window.showInputBox({ + prompt: 'Enter an OpenAI API key', + password: true, + ignoreFocusOut: true, + validateInput: value => (value.trim() ? null : 'API key cannot be empty'), + }); + if (apiKey === undefined) { + return; + } + await aiCore.setOpenAIApiKey(apiKey); + vscode.window.showInformationMessage('OpenAI API key stored securely'); + }) + ); + + context.subscriptions.push( + vscode.commands.registerCommand('openqc.aiClearApiKey', async () => { + await aiCore.clearOpenAIApiKey(); + vscode.window.showInformationMessage('OpenAI API key cleared'); + }) + ); + context.subscriptions.push( vscode.workspace.onDidChangeConfiguration(event => { if (event.affectsConfiguration('openqc.ai')) { @@ -79,12 +102,12 @@ async function aiOptimizeInput(aiCore: AICore): Promise { { location: vscode.ProgressLocation.Notification, title: 'AI optimizing input file...', - cancellable: false, + cancellable: true, }, - async progress => { + async (progress, token) => { progress.report({ message: 'Analyzing input file...' }); - const response = await aiCore.optimizeInput(inputContent, software); + const response = await aiCore.optimizeInput(inputContent, software, token); if (!response.success) { vscode.window.showErrorMessage(`AI optimization failed: ${response.error}`); @@ -147,12 +170,12 @@ async function aiGenerateInput(aiCore: AICore): Promise { { location: vscode.ProgressLocation.Notification, title: 'AI generating input file...', - cancellable: false, + cancellable: true, }, - async progress => { + async (progress, token) => { progress.report({ message: 'Generating input...' }); - const response = await aiCore.generateInput(description, software, ctx); + const response = await aiCore.generateInput(description, software, ctx, token); if (!response.success) { vscode.window.showErrorMessage(`AI generation failed: ${response.error}`); @@ -197,12 +220,12 @@ async function aiExplainParameters(aiCore: AICore): Promise { { location: vscode.ProgressLocation.Notification, title: 'AI analyzing parameters...', - cancellable: false, + cancellable: true, }, - async progress => { + async (progress, token) => { progress.report({ message: 'Analyzing parameters...' }); - const response = await aiCore.explainParameters(inputContent, software); + const response = await aiCore.explainParameters(inputContent, software, token); if (!response.success) { vscode.window.showErrorMessage(`AI explanation failed: ${response.error}`); @@ -271,12 +294,12 @@ async function aiDebugCalculation(aiCore: AICore): Promise { { location: vscode.ProgressLocation.Notification, title: 'AI debugging calculation...', - cancellable: false, + cancellable: true, }, - async progress => { + async (progress, token) => { progress.report({ message: 'Analyzing input and output...' }); - const response = await aiCore.debugCalculation(inputContent, outputContent, software); + const response = await aiCore.debugCalculation(inputContent, outputContent, software, token); if (!response.success) { vscode.window.showErrorMessage(`AI debugging failed: ${response.error}`); @@ -306,6 +329,14 @@ async function aiSettings(aiCore: AICore): Promise { label: '$(symbol-event) Test AI Connection', description: 'Test connection to AI backend', }, + { + label: '$(key) Set OpenAI API Key', + description: 'Store the key in VS Code SecretStorage', + }, + { + label: '$(trash) Clear OpenAI API Key', + description: 'Delete the key from VS Code SecretStorage', + }, ]; const selected = await vscode.window.showQuickPick(items, { @@ -318,13 +349,21 @@ async function aiSettings(aiCore: AICore): Promise { if (selected.label.includes('Configure')) { vscode.commands.executeCommand('workbench.action.openSettings', 'openqc.ai'); + } else if (selected.label.includes('Set OpenAI')) { + vscode.commands.executeCommand('openqc.aiSetApiKey'); + } else if (selected.label.includes('Clear OpenAI')) { + vscode.commands.executeCommand('openqc.aiClearApiKey'); } else if (selected.label.includes('Test')) { const available = await aiCore.isAvailable(); if (available) { vscode.window.showInformationMessage('AI backend is available and ready'); } else { + const errors = [...validation.errors]; + if (config.provider === AIProvider.OpenAI && !(await aiCore.hasOpenAIApiKey())) { + errors.push('OpenAI API key is not configured'); + } vscode.window.showErrorMessage( - `AI backend is not available. ${validation.errors.join(', ')}` + `AI backend is not available${errors.length ? `: ${errors.join(', ')}` : ''}` ); } } diff --git a/tests/python/test_ai_client.py b/tests/python/test_ai_client.py new file mode 100644 index 0000000..8379fe7 --- /dev/null +++ b/tests/python/test_ai_client.py @@ -0,0 +1,175 @@ +import json +import os +import socketserver +import subprocess +import sys +import threading +import time +import unittest +from http.server import BaseHTTPRequestHandler +from pathlib import Path +from unittest.mock import patch + +from openqc.ai.client import AIClient, AIRequest + +PYTHON_ROOT = Path(__file__).resolve().parents[2] / "python" + + +class _ThreadingServer(socketserver.ThreadingMixIn, socketserver.TCPServer): + allow_reuse_address = True + daemon_threads = True + + +class _ProviderHandler(BaseHTTPRequestHandler): + requests = [] + delay = 0.0 + status = 200 + response = {"output_text": "fake response"} + + def do_POST(self): + if self.delay: + time.sleep(self.delay) + length = int(self.headers.get("Content-Length", "0")) + body = json.loads(self.rfile.read(length)) + self.__class__.requests.append((self.path, dict(self.headers), body)) + payload = json.dumps(self.response).encode() + self.send_response(self.status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + try: + self.wfile.write(payload) + except BrokenPipeError: + pass + + def log_message(self, _format, *_args): + return + + +class AIClientFakeHttpTest(unittest.TestCase): + def setUp(self): + _ProviderHandler.requests = [] + _ProviderHandler.delay = 0.0 + _ProviderHandler.status = 200 + _ProviderHandler.response = {"output_text": "fake response"} + self.server = _ThreadingServer(("127.0.0.1", 0), _ProviderHandler) + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() + self.base_url = f"http://127.0.0.1:{self.server.server_address[1]}" + + def tearDown(self): + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=2) + + def _environment(self, **overrides): + values = { + "OPENQC_AI_PROVIDER": "openai", + "OPENQC_AI_MODEL": "gpt-test", + "OPENQC_AI_API_KEY": "test-only-secret", + "OPENQC_AI_OPENAI_URL": self.base_url, + "OPENQC_AI_TIMEOUT_SECONDS": "1", + "OPENQC_AI_MAX_OUTPUT_CHARS": "1000", + } + values.update(overrides) + return patch.dict(os.environ, values, clear=False) + + def test_openai_uses_responses_api_and_bearer_auth(self): + with self._environment(): + response = AIClient().explain(AIRequest("explain", "ENCUT=500", "vasp")) + + self.assertTrue(response.success) + self.assertEqual(response.content, "fake response") + path, headers, body = _ProviderHandler.requests[0] + self.assertEqual(path, "/responses") + self.assertEqual(headers["Authorization"], "Bearer test-only-secret") + self.assertEqual(body["model"], "gpt-test") + self.assertEqual(body["max_output_tokens"], 2048) + self.assertIn("ENCUT=500", body["input"]) + + def test_module_entrypoint_succeeds_against_fake_http_server(self): + env = os.environ.copy() + env.update( + { + "PYTHONPATH": str(PYTHON_ROOT), + "OPENQC_AI_PROVIDER": "openai", + "OPENQC_AI_MODEL": "gpt-test", + "OPENQC_AI_API_KEY": "test-only-secret", + "OPENQC_AI_OPENAI_URL": self.base_url, + } + ) + completed = subprocess.run( + [sys.executable, "-m", "openqc.ai.client", "explain"], + input=json.dumps( + {"type": "explain", "content": "input", "software": "orca"} + ), + text=True, + capture_output=True, + env=env, + cwd=PYTHON_ROOT, + timeout=5, + check=True, + ) + + payload = json.loads(completed.stdout) + self.assertTrue(payload["success"]) + self.assertEqual(payload["content"], "fake response") + self.assertNotIn("test-only-secret", completed.stdout + completed.stderr) + self.assertNotIn("RuntimeWarning", completed.stderr) + + def test_package_exports_remain_available_with_lazy_client_import(self): + from openqc.ai import AIClient as ExportedClient + + self.assertIs(ExportedClient, AIClient) + + def test_missing_openai_credential_fails_without_http(self): + with self._environment(OPENQC_AI_API_KEY=""): + response = AIClient().explain(AIRequest("explain", "input", "vasp")) + + self.assertFalse(response.success) + self.assertEqual(response.error, "OpenAI API key is not configured") + self.assertEqual(_ProviderHandler.requests, []) + + def test_provider_error_is_redacted(self): + _ProviderHandler.status = 401 + _ProviderHandler.response = { + "error": {"message": "bad Bearer test-only-secret"} + } + with self._environment(): + response = AIClient().explain(AIRequest("explain", "input", "vasp")) + + self.assertFalse(response.success) + self.assertIn("HTTP 401", response.error) + self.assertNotIn("test-only-secret", response.error) + + def test_timeout_is_bounded(self): + _ProviderHandler.delay = 0.2 + with self._environment(OPENQC_AI_TIMEOUT_SECONDS="0.05"): + response = AIClient().explain(AIRequest("explain", "input", "vasp")) + + self.assertFalse(response.success) + self.assertIn("timed out", response.error) + + def test_output_is_truncated_to_configured_limit(self): + _ProviderHandler.response = {"output_text": "x" * 1000} + with self._environment(OPENQC_AI_MAX_OUTPUT_CHARS="256"): + response = AIClient().explain(AIRequest("explain", "input", "vasp")) + + self.assertTrue(response.success) + self.assertEqual(len(response.content), 256) + self.assertEqual(response.metadata, {"truncated": True}) + + def test_ollama_generate_api(self): + _ProviderHandler.response = {"response": "local response"} + with self._environment( + OPENQC_AI_PROVIDER="ollama", OPENQC_AI_OLLAMA_URL=self.base_url + ): + response = AIClient().explain(AIRequest("explain", "input", "vasp")) + + self.assertTrue(response.success) + self.assertEqual(response.content, "local response") + self.assertEqual(_ProviderHandler.requests[0][0], "/api/generate") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/python/test_mcp_contract.py b/tests/python/test_mcp_contract.py new file mode 100644 index 0000000..f8b6ca2 --- /dev/null +++ b/tests/python/test_mcp_contract.py @@ -0,0 +1,112 @@ +import json +import unittest +from unittest.mock import patch + +from openqc.mcp_server import _HANDLERS, TOOLS, handle_request + + +def request(request_id, method, params=None): + value = {"jsonrpc": "2.0", "id": request_id, "method": method} + if params is not None: + value["params"] = params + return value + + +class McpContractTest(unittest.TestCase): + def test_tools_list_and_every_declared_tool_are_callable(self): + listed = handle_request(request("list-id", "tools/list")) + tools = listed["result"]["tools"] + self.assertEqual(tools, TOOLS) + self.assertEqual(len(tools), 5) + + fake_handlers = { + tool["name"]: (lambda args, name=tool["name"]: {"tool": name, "args": args}) + for tool in tools + } + arguments = { + "openqc.parseStructure": {"path": "structure.xyz"}, + "openqc.parseCalculationOutput": {"path": "calculation.out"}, + "openqc.checkPythonBackend": {}, + "openqc.generateSupercell": {"path": "POSCAR"}, + "openqc.summarizeDataset": {"path": "dataset"}, + } + with patch.dict(_HANDLERS, fake_handlers, clear=True): + for index, tool in enumerate(tools): + request_id = f"tool-{index}" + response = handle_request( + request( + request_id, + "tools/call", + {"name": tool["name"], "arguments": arguments[tool["name"]]}, + ) + ) + self.assertEqual(response["id"], request_id) + payload = json.loads(response["result"]["content"][0]["text"]) + self.assertEqual(payload["tool"], tool["name"]) + + self.assertEqual( + json.loads(response["result"]["content"][0]["text"])["args"]["format"], + "auto", + ) + + def test_invalid_jsonrpc_and_params_have_stable_codes_and_ids(self): + invalid = handle_request({"jsonrpc": "1.0", "id": 0, "method": "tools/list"}) + self.assertEqual(invalid["id"], 0) + self.assertEqual(invalid["error"]["code"], -32600) + + missing = handle_request( + request( + 42, + "tools/call", + {"name": "openqc.parseStructure", "arguments": {}}, + ) + ) + self.assertEqual(missing["id"], 42) + self.assertEqual(missing["error"]["code"], -32602) + + wrong_type = handle_request( + request( + "same-id", + "tools/call", + { + "name": "openqc.generateSupercell", + "arguments": {"path": "POSCAR", "na": 1.5}, + }, + ) + ) + self.assertEqual(wrong_type["id"], "same-id") + self.assertEqual(wrong_type["error"]["code"], -32602) + + def test_optional_dependency_and_internal_errors_are_normalized(self): + def missing(_arguments): + raise ImportError("install dpdata") + + def broken(_arguments): + raise RuntimeError("Bearer example-credential") + + call = lambda request_id: request( + request_id, + "tools/call", + {"name": "openqc.summarizeDataset", "arguments": {"path": "dataset"}}, + ) + + with patch.dict(_HANDLERS, {"openqc.summarizeDataset": missing}): + dependency = handle_request(call("dependency-id")) + self.assertEqual(dependency["id"], "dependency-id") + self.assertEqual(dependency["error"]["code"], -32001) + + with patch.dict(_HANDLERS, {"openqc.summarizeDataset": broken}): + internal = handle_request(call("internal-id")) + self.assertEqual(internal["id"], "internal-id") + self.assertEqual(internal["error"]["code"], -32603) + self.assertNotIn("example-credential", internal["error"]["message"]) + + def test_notifications_do_not_receive_responses(self): + self.assertIsNone( + handle_request({"jsonrpc": "2.0", "method": "notifications/initialized"}) + ) + self.assertIsNone(handle_request({"jsonrpc": "2.0", "method": "unknown"})) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/ai/AICore.test.ts b/tests/unit/ai/AICore.test.ts index 841a4f0..bccfbc6 100644 --- a/tests/unit/ai/AICore.test.ts +++ b/tests/unit/ai/AICore.test.ts @@ -1,375 +1,268 @@ -/** - * AICore Unit Tests - */ +import { AICore, AICoreFactory, AIProvider, OPENAI_API_KEY_SECRET } from '../../../src/ai/AICore'; -import { AICore, AICoreFactory, AIProvider, AIRequestType } from '../../../src/ai/AICore'; - -// Mock vscode const mockGetConfiguration = jest.fn(); -const mockContext = { - extensionPath: '/test/extension', -} as any; +const mockSpawn = jest.fn(); +const secretGet = jest.fn(); +const secretStore = jest.fn(); +const secretDelete = jest.fn(); jest.mock('vscode', () => ({ - workspace: { - getConfiguration: () => mockGetConfiguration(), - }, - window: { - showErrorMessage: jest.fn(), - showInformationMessage: jest.fn(), - }, + workspace: { getConfiguration: () => mockGetConfiguration() }, })); -// Mock child_process -const mockSpawn = jest.fn(); jest.mock('child_process', () => ({ - spawn: (...args: any[]) => mockSpawn(...args), + spawn: (...args: unknown[]) => mockSpawn(...args), })); -describe('AICore', () => { - let aiCore: AICore; +function config(overrides: Record = {}): void { + const values: Record = { + enabled: true, + provider: AIProvider.Ollama, + model: 'llama3.2', + openaiBaseUrl: 'https://api.openai.invalid/v1', + ollamaUrl: 'http://localhost:11434', + maxTokens: 128, + maxOutputChars: 1024, + timeoutSeconds: 5, + temperature: 0.2, + pythonPath: 'python-test', + ...overrides, + }; + mockGetConfiguration.mockReturnValue({ + get: jest.fn((key: string, defaultValue?: unknown) => values[key] ?? defaultValue), + }); +} + +function childProcess( + options: { + stdout?: string; + stderr?: string; + code?: number; + neverClose?: boolean; + } = {} +) { + const kill = jest.fn(); + return { + stdin: { write: jest.fn(), end: jest.fn() }, + stdout: { + on: jest.fn((event: string, callback: (data: string) => void) => { + if (event === 'data' && options.stdout !== undefined) { + callback(options.stdout); + } + }), + }, + stderr: { + on: jest.fn((event: string, callback: (data: string) => void) => { + if (event === 'data' && options.stderr !== undefined) { + callback(options.stderr); + } + }), + }, + on: jest.fn((event: string, callback: (value: number) => void) => { + if (event === 'close' && !options.neverClose) { + callback(options.code ?? 0); + } + }), + kill, + }; +} + +describe('AICore secure bridge', () => { + const context = { + extensionPath: '/extension', + secrets: { get: secretGet, store: secretStore, delete: secretDelete }, + } as any; beforeEach(() => { jest.clearAllMocks(); + jest.useRealTimers(); AICoreFactory.reset(); - - // Default mock configuration - mockGetConfiguration.mockReturnValue({ - get: jest.fn((key: string, defaultValue?: any) => { - const config: Record = { - enabled: true, - provider: AIProvider.Ollama, - apiKey: undefined, - model: 'llama2', - ollamaUrl: 'http://localhost:11434', - maxTokens: 2048, - temperature: 0.7, - pythonPath: 'python3', - }; - return config[key] ?? defaultValue; - }), - }); - - aiCore = new AICore(mockContext); + config(); + secretGet.mockResolvedValue(undefined); + secretStore.mockResolvedValue(undefined); + secretDelete.mockResolvedValue(undefined); }); - describe('Configuration', () => { - it('should load configuration correctly', () => { - const config = aiCore.getConfig(); - - expect(config.enabled).toBe(true); - expect(config.provider).toBe(AIProvider.Ollama); - expect(config.model).toBe('llama2'); - expect(config.ollamaUrl).toBe('http://localhost:11434'); - expect(config.maxTokens).toBe(2048); - expect(config.temperature).toBe(0.7); - }); - - it('should refresh configuration', () => { - mockGetConfiguration.mockReturnValue({ - get: jest.fn((key: string, defaultValue?: any) => { - const config: Record = { - enabled: false, - provider: AIProvider.OpenAI, - apiKey: 'test-key', - model: 'gpt-4', - ollamaUrl: 'http://localhost:11434', - maxTokens: 4096, - temperature: 0.5, - pythonPath: 'python3', - }; - return config[key] ?? defaultValue; - }), - }); - - aiCore.refreshConfig(); - const config = aiCore.getConfig(); - - expect(config.enabled).toBe(false); - expect(config.provider).toBe(AIProvider.OpenAI); - expect(config.apiKey).toBe('test-key'); - expect(config.model).toBe('gpt-4'); - }); - - it('should check if AI is enabled', () => { - expect(aiCore.isEnabled()).toBe(true); + it('stores and clears the OpenAI credential only through SecretStorage', async () => { + const core = new AICore(context); + await core.setOpenAIApiKey(' test-secret '); + await core.clearOpenAIApiKey(); - mockGetConfiguration.mockReturnValue({ - get: jest.fn(() => false), - }); - aiCore.refreshConfig(); + expect(secretStore).toHaveBeenCalledWith(OPENAI_API_KEY_SECRET, 'test-secret'); + expect(secretDelete).toHaveBeenCalledWith(OPENAI_API_KEY_SECRET); + expect(core.getConfig()).not.toHaveProperty('apiKey'); + }); - expect(aiCore.isEnabled()).toBe(false); + it('refreshes settings and validates unsafe bounds', () => { + const core = new AICore(context); + expect(core.isEnabled()).toBe(true); + config({ enabled: false, model: '', temperature: 3, timeoutSeconds: 0, maxOutputChars: 10 }); + + core.refreshConfig(); + + expect(core.isEnabled()).toBe(false); + expect(core.validateConfig()).toEqual({ + valid: false, + errors: [ + 'AI features are disabled', + 'Model name is required', + 'Temperature must be between 0 and 2', + 'Timeout must be between 0 and 600 seconds', + 'Maximum output characters must be between 256 and 1000000', + ], }); }); - describe('Validation', () => { - it('should validate configuration when enabled', () => { - const validation = aiCore.validateConfig(); + it('reports whether SecretStorage contains an OpenAI key', async () => { + secretGet.mockResolvedValueOnce('credential').mockResolvedValueOnce(undefined); + const core = new AICore(context); - expect(validation.valid).toBe(true); - expect(validation.errors).toHaveLength(0); - }); - - it('should fail validation when disabled', () => { - mockGetConfiguration.mockReturnValue({ - get: jest.fn((key: string) => { - if (key === 'enabled') return false; - return undefined; - }), - }); - aiCore.refreshConfig(); + await expect(core.hasOpenAIApiKey()).resolves.toBe(true); + await expect(core.hasOpenAIApiKey()).resolves.toBe(false); + await expect(core.setOpenAIApiKey(' ')).rejects.toThrow('cannot be empty'); + }); - const validation = aiCore.validateConfig(); + it('rejects OpenAI requests with a missing SecretStorage credential', async () => { + config({ provider: AIProvider.OpenAI, model: 'gpt-test' }); + const result = await new AICore(context).explainParameters('input', 'vasp'); - expect(validation.valid).toBe(false); - expect(validation.errors).toContain('AI features are disabled'); - }); - - it('should require API key for OpenAI', () => { - mockGetConfiguration.mockReturnValue({ - get: jest.fn((key: string, defaultValue?: any) => { - const config: Record = { - enabled: true, - provider: AIProvider.OpenAI, - apiKey: undefined, - model: 'gpt-4', - ollamaUrl: 'http://localhost:11434', - maxTokens: 2048, - temperature: 0.7, - }; - return config[key] ?? defaultValue; - }), - }); - aiCore.refreshConfig(); - - const validation = aiCore.validateConfig(); - - expect(validation.valid).toBe(false); - expect(validation.errors).toContain('OpenAI API key is required'); - }); + expect(result).toEqual({ success: false, error: 'OpenAI API key is not configured' }); + expect(mockSpawn).not.toHaveBeenCalled(); + }); - it('should validate temperature range', () => { - mockGetConfiguration.mockReturnValue({ - get: jest.fn((key: string, defaultValue?: any) => { - const config: Record = { - enabled: true, - provider: AIProvider.Ollama, - model: 'llama2', - temperature: 3.0, // Invalid: > 2 - }; - return config[key] ?? defaultValue; - }), - }); - aiCore.refreshConfig(); - - const validation = aiCore.validateConfig(); - - expect(validation.valid).toBe(false); - expect(validation.errors).toContain('Temperature must be between 0 and 2'); - }); + it('runs the Python client as a module and passes the secret only in the child environment', async () => { + config({ provider: AIProvider.OpenAI, model: 'gpt-test' }); + secretGet.mockResolvedValue('secret-value'); + mockSpawn.mockReturnValue( + childProcess({ stdout: JSON.stringify({ success: true, content: 'ok' }) }) + ); + + const result = await new AICore(context).explainParameters('input', 'vasp'); + + expect(result).toMatchObject({ success: true, content: 'ok' }); + const [executable, args, options] = mockSpawn.mock.calls[0]; + expect(executable).toBe('python-test'); + expect(args).toEqual(['-m', 'openqc.ai.client', 'explain']); + expect(options.env.OPENQC_AI_API_KEY).toBe('secret-value'); + expect(JSON.stringify(args)).not.toContain('secret-value'); }); - describe('AI Operations', () => { - const mockSuccessResponse = { + it('routes generate, debug, and availability operations to their module commands', async () => { + mockSpawn.mockImplementation(() => + childProcess({ stdout: JSON.stringify({ success: true, generatedInput: 'input' }) }) + ); + const core = new AICore(context); + + await expect(core.generateInput('water', 'cp2k', { structure: 'xyz' })).resolves.toMatchObject({ success: true, - content: 'Test response', - suggestions: [], - }; - - beforeEach(() => { - // Mock successful spawn - mockSpawn.mockReturnValue({ - stdin: { write: jest.fn(), end: jest.fn() }, - stdout: { - on: jest.fn((event: string, callback: Function) => { - if (event === 'data') { - callback(JSON.stringify(mockSuccessResponse)); - } - }), - }, - stderr: { on: jest.fn() }, - on: jest.fn((event: string, callback: Function) => { - if (event === 'close') { - callback(0); - } - }), - }); }); - - it('should optimize input', async () => { - const result = await aiCore.optimizeInput('test input', 'vasp'); - - expect(result.success).toBe(true); - expect(mockSpawn).toHaveBeenCalled(); + await expect(core.debugCalculation('input', 'output', 'orca')).resolves.toMatchObject({ + success: true, }); + await expect(core.isAvailable()).resolves.toBe(true); - it('should generate input', async () => { - const result = await aiCore.generateInput('test description', 'cp2k'); - - expect(result.success).toBe(true); - expect(mockSpawn).toHaveBeenCalled(); - }); + expect(mockSpawn.mock.calls.map(call => call[1][2])).toEqual(['generate', 'debug', 'check']); + }); - it('should explain parameters', async () => { - const result = await aiCore.explainParameters('test input', 'gaussian'); + it('does not check availability while AI is disabled', async () => { + config({ enabled: false }); + await expect(new AICore(context).isAvailable()).resolves.toBe(false); + expect(mockSpawn).not.toHaveBeenCalled(); + }); - expect(result.success).toBe(true); - expect(mockSpawn).toHaveBeenCalled(); + it('short-circuits an already-cancelled request', async () => { + const token = { isCancellationRequested: true } as any; + await expect(new AICore(context).optimizeInput('input', 'vasp', token)).resolves.toEqual({ + success: false, + error: 'AI request was cancelled', }); + expect(mockSpawn).not.toHaveBeenCalled(); + }); - it('should debug calculation', async () => { - const result = await aiCore.debugCalculation('input', 'output', 'qe'); - - expect(result.success).toBe(true); - expect(mockSpawn).toHaveBeenCalled(); - }); + it('kills a timed-out bridge process', async () => { + jest.useFakeTimers(); + config({ timeoutSeconds: 0.01 }); + const child = childProcess({ neverClose: true }); + mockSpawn.mockReturnValue(child); - it('should handle Python errors', async () => { - mockSpawn.mockReturnValue({ - stdin: { write: jest.fn(), end: jest.fn() }, - stdout: { on: jest.fn() }, - stderr: { - on: jest.fn((event: string, callback: Function) => { - if (event === 'data') { - callback('Python error'); - } - }), - }, - on: jest.fn((event: string, callback: Function) => { - if (event === 'close') { - callback(1); - } - }), - }); - - const result = await aiCore.optimizeInput('test', 'vasp'); - - expect(result.success).toBe(false); - expect(result.error).toContain('AI backend failed'); - }); + const pending = new AICore(context).optimizeInput('input', 'vasp'); + await jest.advanceTimersByTimeAsync(100); - it('should handle invalid JSON response', async () => { - mockSpawn.mockReturnValue({ - stdin: { write: jest.fn(), end: jest.fn() }, - stdout: { - on: jest.fn((event: string, callback: Function) => { - if (event === 'data') { - callback('invalid json'); - } - }), - }, - stderr: { on: jest.fn() }, - on: jest.fn((event: string, callback: Function) => { - if (event === 'close') { - callback(0); - } - }), - }); - - const result = await aiCore.optimizeInput('test', 'vasp'); - - expect(result.success).toBe(false); - expect(result.error).toContain('Failed to parse AI response'); + await expect(pending).resolves.toEqual({ + success: false, + error: 'AI request timed out after 0.01 seconds', }); + expect(child.kill).toHaveBeenCalled(); }); - describe('Availability Check', () => { - it('should return false when disabled', async () => { - mockGetConfiguration.mockReturnValue({ - get: jest.fn(() => false), - }); - aiCore.refreshConfig(); + it('kills the process when VS Code cancellation is requested', async () => { + let cancel: (() => void) | undefined; + const token = { + isCancellationRequested: false, + onCancellationRequested: jest.fn((callback: () => void) => { + cancel = callback; + return { dispose: jest.fn() }; + }), + } as any; + const child = childProcess({ neverClose: true }); + mockSpawn.mockReturnValue(child); - const available = await aiCore.isAvailable(); - expect(available).toBe(false); - }); + const pending = new AICore(context).optimizeInput('input', 'vasp', token); + await Promise.resolve(); + cancel?.(); - it('should check availability via Python backend', async () => { - mockSpawn.mockReturnValue({ - stdin: { write: jest.fn(), end: jest.fn() }, - stdout: { - on: jest.fn((event: string, callback: Function) => { - if (event === 'data') { - callback(JSON.stringify({ success: true })); - } - }), - }, - stderr: { on: jest.fn() }, - on: jest.fn((event: string, callback: Function) => { - if (event === 'close') { - callback(0); - } - }), - }); - - const available = await aiCore.isAvailable(); - expect(available).toBe(true); - }); + await expect(pending).resolves.toEqual({ success: false, error: 'AI request was cancelled' }); + expect(child.kill).toHaveBeenCalled(); }); - describe('Available Models', () => { - it('should return OpenAI models when provider is OpenAI', async () => { - // Mock OpenAI provider - mockGetConfiguration.mockReturnValue({ - get: jest.fn((key: string, defaultValue?: any) => { - const config: Record = { - enabled: true, - provider: AIProvider.OpenAI, - apiKey: 'test-key', - model: 'gpt-4', - ollamaUrl: 'http://localhost:11434', - maxTokens: 2048, - temperature: 0.7, - pythonPath: 'python3', - }; - return config[key] ?? defaultValue; - }), - }); - aiCore.refreshConfig(); - - const models = await aiCore.getAvailableModels(); - - expect(models).toContain('gpt-4'); - expect(models).toContain('gpt-3.5-turbo'); - }); - - it('should return Ollama models', async () => { - global.fetch = jest.fn().mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ models: [{ name: 'llama2' }, { name: 'mistral' }] }), - }); + it('bounds stdout and does not echo invalid output', async () => { + config({ maxOutputChars: 256 }); + const child = childProcess({ stdout: 'x'.repeat(5000), neverClose: true }); + mockSpawn.mockReturnValue(child); - const models = await aiCore.getAvailableModels(); + const result = await new AICore(context).optimizeInput('input', 'vasp'); - expect(models).toContain('llama2'); - expect(models).toContain('mistral'); - }); + expect(result.error).toBe('AI backend output exceeded the configured safe limit'); + expect(result).not.toHaveProperty('metadata.raw_output'); + expect(child.kill).toHaveBeenCalled(); + }); - it('should handle Ollama fetch error', async () => { - global.fetch = jest.fn().mockRejectedValue(new Error('Connection failed')); + it('redacts the credential from child-process errors', async () => { + config({ provider: AIProvider.OpenAI }); + secretGet.mockResolvedValue('example-credential'); + mockSpawn.mockReturnValue( + childProcess({ code: 1, stderr: 'Authorization: Bearer example-credential' }) + ); - const models = await aiCore.getAvailableModels(); + const result = await new AICore(context).optimizeInput('input', 'vasp'); - expect(models).toContain('llama2'); - expect(models).toContain('codellama'); - }); + expect(result.error).toContain('[REDACTED]'); + expect(result.error).not.toContain('example-credential'); }); - describe('Factory', () => { - it('should return singleton instance', () => { - const instance1 = AICoreFactory.getInstance(mockContext); - const instance2 = AICoreFactory.getInstance(mockContext); + it('returns configured OpenAI model and discovers Ollama models', async () => { + config({ provider: AIProvider.OpenAI, model: 'gpt-test' }); + await expect(new AICore(context).getAvailableModels()).resolves.toEqual(['gpt-test']); - expect(instance1).toBe(instance2); - }); - - it('should reset instance', () => { - const instance1 = AICoreFactory.getInstance(mockContext); - AICoreFactory.reset(); - const instance2 = AICoreFactory.getInstance(mockContext); + config({ provider: AIProvider.Ollama }); + const fetchMock = jest + .spyOn(global, 'fetch') + .mockResolvedValueOnce({ + ok: true, + json: async () => ({ models: [{ name: 'local-a' }, {}, { name: 'local-b' }] }), + } as any) + .mockRejectedValueOnce(new Error('offline')); + const core = new AICore(context); + + await expect(core.getAvailableModels()).resolves.toEqual(['local-a', 'local-b']); + await expect(core.getAvailableModels()).resolves.toEqual([]); + fetchMock.mockRestore(); + }); - expect(instance1).not.toBe(instance2); - }); + it('reuses and resets the factory singleton', () => { + const first = AICoreFactory.getInstance(context); + expect(AICoreFactory.getInstance(context)).toBe(first); + AICoreFactory.reset(); + expect(AICoreFactory.getInstance(context)).not.toBe(first); }); }); diff --git a/tests/unit/ai/aiSecurityContract.test.ts b/tests/unit/ai/aiSecurityContract.test.ts new file mode 100644 index 0000000..d630d3d --- /dev/null +++ b/tests/unit/ai/aiSecurityContract.test.ts @@ -0,0 +1,20 @@ +import * as fs from 'fs'; +import * as path from 'path'; + +describe('AI manifest security contract', () => { + const root = path.resolve(__dirname, '../../..'); + const manifest = JSON.parse(fs.readFileSync(path.join(root, 'package.json'), 'utf8')); + + it('does not contribute a plaintext API-key setting', () => { + const properties = manifest.contributes.configuration.properties; + expect(properties).not.toHaveProperty('openqc.ai.apiKey'); + }); + + it('contributes and activates SecretStorage set and clear commands', () => { + const commands = manifest.contributes.commands.map((item: { command: string }) => item.command); + for (const command of ['openqc.aiSetApiKey', 'openqc.aiClearApiKey']) { + expect(commands).toContain(command); + expect(manifest.activationEvents).toContain(`onCommand:${command}`); + } + }); +}); From 84d657fff9b683f595cb325eb9faf1e45cfce86a Mon Sep 17 00:00:00 2001 From: OpenQC Test Date: Thu, 16 Jul 2026 16:52:58 +0800 Subject: [PATCH 2/2] docs: record PR contract test heading --- docs/agent-memory.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/agent-memory.md b/docs/agent-memory.md index d280f99..d7b634a 100644 --- a/docs/agent-memory.md +++ b/docs/agent-memory.md @@ -8,6 +8,8 @@ Fake local HTTP servers cover provider behavior; tests must never send live provider requests. - The dependency-free MCP server declares five tools. Keep the tool schemas, handler routing, JSON-RPC validation, request IDs, and normalized optional-dependency errors in one module. +- The PR contract recognizes test evidence only when the body contains `Tests`, `Test Plan`, or + `TDD Evidence`; a generic `Validation` heading does not satisfy the check. - `make lint` is currently blocked by 544 pre-existing Ruff errors under `core/`. `make test` reaches a missing `pytest` executable and a nonexistent `core/tests/unit` path after the Jest unit suite passes; use the maintained root Python tests plus npm/Jest gates for this slice.