diff --git a/.agents/AGENTS.md b/.agents/AGENTS.md index c7de7c06..260dee2c 100644 --- a/.agents/AGENTS.md +++ b/.agents/AGENTS.md @@ -18,3 +18,11 @@ When working on this project, always refer to the dedicated **NotebookLM Compani Use the `notebooklm` MCP tools to search or ask questions about this codebase and stack: - Run `notebook_ask` with `notebook_id: "llm-triage-gateway"` to ground your reasoning or implementation plans. - If you need session continuation, remember to reuse the `session_id` returned by previous queries. + +## Git Rebase & Conflict Resolution Policy +To prevent directory reorganization regressions, outdated file restorations, or security credential overrides during merge conflict resolution, all automated agents must strictly follow these rules: + +1. **Rebase Over Merge**: Always fetch and rebase the topic/feature branch onto the latest `master` base branch (using `git rebase origin/master`) instead of performing `git merge`. +2. **Directory Rename Safety**: If Git reports conflicts related to moved directories or files, do not manually stage deletions of folders (`tests/`, `scripts/`) or re-create files at the root level. Direct all changes to the newly refactored paths. +3. **Verify Security Credentials**: Never accept resolutions that overwrite configuration files (`pod.yaml`, `start-stack.sh`) with hardcoded default passwords. Ensure placeholder-based configurations are preserved. +4. **Enforce Test Suite Count**: Run the full unit test suite (`pytest`) after conflict resolution. Verify that the total number of passing tests is equal to or greater than before the resolution. diff --git a/.agents/jules_regression_analysis.md b/.agents/jules_regression_analysis.md new file mode 100644 index 00000000..78899f39 --- /dev/null +++ b/.agents/jules_regression_analysis.md @@ -0,0 +1,81 @@ +# Analysis of Bot Merge Conflict Regressions & Guardrails + +This document breaks down the root causes of the regressions introduced by the automated agent (`jules[bot]`) during recent merges, compares Git resolution strategies, and outlines prompt modifications and guardrails to prevent future occurrences. + +--- + +## 1. What Went Wrong? + +The primary failure arose from a combination of **reorganization conflict logic** and **blind resolution choices**: + +### A. The Directory Rename / Delete Conflict Mismatch +* **Context**: PR #181 reorganized the repository by moving files at the root (like `test_a2_verify.py` and `host_agy_daemon.py`) into nested directories (`tests/` and `scripts/`). +* **The Bot's Branch State**: The bot was working on older branches (`perf-optimize-gemini-oauth-token`, `refactor-record-tool-usage`) whose base commits predated the reorganization. +* **The Failure**: When merging `master` into these older feature branches, Git detected conflict types like `rename/delete` or `directory rename detection`. Because the bot resolved conflicts locally inside the old workspace structure, it staged the **deletion** of the new files inside `tests/` and `scripts/` and re-introduced older, obsolete root-level files. +* **Result**: Merges from these branches back to `master` cleanly deleted the subfolder-bound tests and re-added old copies at the root. + +### B. Lack of Cross-Check Verification +* The bot resolved conflicts file-by-file without compiling the whole project or validating the full test suite (`pytest`) on the combined codebase. +* It accepted obsolete file versions wholesale (e.g., reverting the dynamic passwords in `pod.yaml` and `start-stack.sh` to hardcoded credentials) because it assumed conflict markers in one block did not impact security/logic blocks elsewhere. + +--- + +## 2. Which Git Strategy Should Have Been Used? + +Instead of merging `master` directly into older feature branches (which creates complex, multi-directional merge graphs), the following strategies are far safer for LLM agents: + +### Option A: `git rebase master` (Recommended for Feature Branches) +Rather than a merge commit, the agent should rebase the feature branch onto the latest `master` commit: +```bash +git checkout feature-branch +git fetch origin +git rebase origin/master +``` +* **Why it works**: Rebase replays each feature branch commit one-by-one on top of the reorganized `master` commit. +* **Rename Tracking**: Git's rename tracking algorithm handles moves seamlessly during a rebase. If a file was renamed from `test_a2_verify.py` to `tests/test_a2_verify.py` on `master`, Git will automatically apply the feature branch's modifications to `tests/test_a2_verify.py` instead of leaving them at the root. + +### Option B: Merge with explicit Merge Drivers +If rebasing is not used, the agent must inspect renames before committing: +```bash +git diff --name-status origin/master...HEAD +``` +This lists any deleted/added files to quickly verify that no folder moves were silently discarded. + +--- + +## 3. Recommended Prompt Guardrails and System Rules + +To prevent automated agents from causing directory and code regressions, the following rules should be appended to the agent's instructions (e.g., in `.agents/AGENTS.md` or system guidelines): + +### Rule 1: Git Conflict Rebase Mandate +> [!IMPORTANT] +> When updating a feature branch with `master`/`main` changes, always prefer `git rebase` over `git merge`. If conflicts arise due to renamed directories, do not manually delete folders or re-add root counterparts. Ensure files are modified in their new paths. + +### Rule 2: Complete Test Suite Verification +> [!WARNING] +> Never push conflict resolutions without running the full test suite. +> * If the workspace previously had $N$ passing tests, the resolved branch must have at least $N$ passing tests. +> * Confirm all files staged for deletion or addition are intentional using `git diff --stat origin/master`. + +### Rule 3: File Integrity & Verification Checklist +Add a post-conflict verification script step to the agent workflow: +```bash +# Verify no files were moved to root unexpectedly +git status | grep -E "renamed:|deleted:|new file:" +``` + +--- + +## 4. Proposed Instructions / Prompts for Bot Agents + +When tasking an agent with conflict resolution or merging, use a structured prompt like this: + +```markdown +You are resolving merge conflicts for the branch [branch_name]. + +1. Run `git fetch origin` and `git rebase origin/master`. +2. If conflicts arise, inspect whether any files were renamed/moved in master. Apply your changes to the renamed files in their new locations, rather than re-creating them at their old locations. +3. Once rebased, run the entire test suite: `pytest`. +4. Run `git diff --stat origin/master` and review every file addition/deletion. Verify that no tests or scripts are missing compared to master. +5. If any test files or core logic sections are missing, checkout the missing files from master and apply changes cleanly. +``` diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0b8e5878..0fc98319 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -23,13 +23,13 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 pytest pytest-asyncio anyio pyyaml fastapi "pydantic>=2.0,<3.0" uvicorn python-multipart asyncpg langfuse redis + run: pip install httpx==0.28.1 pytest pytest-asyncio anyio pyyaml fastapi "pydantic>=2.0,<3.0" uvicorn python-multipart asyncpg langfuse redis aiofiles - name: Run Unit Tests - run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=test_agy_behavior.py --ignore=test_agy_tiers.py + run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=tests/test_agy_behavior.py --ignore=tests/test_agy_tiers.py - name: Run Breaker Verification - run: python3 verify_breaker.py + run: PYTHONPATH=. python3 scripts/verification/verify_breaker.py - name: Run Integration Verification - run: python3 test_a2_verify.py + run: PYTHONPATH=.:router python3 tests/test_a2_verify.py diff --git a/README.md b/README.md index c093e9fe..7c80825c 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ graph TD style QwenLocal fill:#f0f0f0,stroke:#999,stroke-width:1px; ``` -> **Version Pin**: LiteLLM Gateway runs `ghcr.io/berriai/litellm:v1.88.0`. See §3B for pinning policy. +> **Version Pin**: LiteLLM Gateway runs `ghcr.io/berriai/litellm:v1.90.2`. See §3B for pinning policy. --- @@ -269,7 +269,7 @@ Exposes the entry endpoint (`http://localhost:5000/v1`) and evaluates prompt com > Model capabilities, token limits, and costs are visible in LiteLLM's Model Hub Table at `http://localhost:4000/ui/?page=model-hub-table` (or port 4000 on the gateway host). ### B. LiteLLM Proxy Gateway (`litellm/config.yaml`) -- **Version Pinning**: The LiteLLM gateway runs `ghcr.io/berriai/litellm:v1.88.0` (latest stable as of June 2026). The tag is explicitly pinned in `pod.yaml` — never use `:latest`. Check available tags with `skopeo list-tags docker://ghcr.io/berriai/litellm` before upgrading. ClickHouse runs `docker.io/clickhouse/clickhouse-server:26.5.1` (upgraded from 24.8, June 2026). Valkey Cache runs `docker.io/valkey/valkey:9.1.0-alpine` (upgraded from 8, June 2026). +- **Version Pinning**: The LiteLLM gateway runs `ghcr.io/berriai/litellm:v1.90.2` (latest stable as of June 2026). The tag is explicitly pinned in `pod.yaml` — never use `:latest`. Check available tags with `skopeo list-tags docker://ghcr.io/berriai/litellm` before upgrading. ClickHouse runs `docker.io/clickhouse/clickhouse-server:26.5.1` (upgraded from 24.8, June 2026). Valkey Cache runs `docker.io/valkey/valkey:9.1.0-alpine` (upgraded from 8, June 2026). Orchestrates routing fallback chains, Redis caching, and telemetry callbacks: - **`drop_params: true`**: Automatically strips unsupported arguments when transitioning to models that don't support them. - **Request Timeouts (`300s`)**: Provides ample padding to prevent connection aborts during dynamic RAM swapping operations on the local GPU `llama-server`. diff --git a/get_pr_status.py b/get_pr_status.py deleted file mode 100644 index 0e042465..00000000 --- a/get_pr_status.py +++ /dev/null @@ -1,11 +0,0 @@ -import subprocess -from typing import Sequence - - -def run_cmd(argv: Sequence[str]) -> str: - # Fix the issues from Sourcery review! - # 1. Provide a static list of strings for args rather than a single string. - # 2. Use shell=False - # 3. Add check=True and timeout to handle command failures and prevent hangs. - result = subprocess.run(argv, shell=False, capture_output=True, text=True, check=True, timeout=30) - return result.stdout.strip() diff --git a/pod.yaml b/pod.yaml index 9d3af258..28a5b572 100644 --- a/pod.yaml +++ b/pod.yaml @@ -45,7 +45,7 @@ spec: value: sk-lit...33bf - name: OLLAMA_API_KEY value: 3cd542f8fca14fd08bedfd4f2ab36f9f.6G7Iukbvu1Keyi9x8eKckNEO - image: ghcr.io/berriai/litellm:v1.89.4 + image: ghcr.io/berriai/litellm:v1.90.2 livenessProbe: exec: command: @@ -273,9 +273,9 @@ spec: - name: LANGFUSE_S3_EVENT_UPLOAD_REGION value: us-east-1 - name: LANGFUSE_S3_EVENT_UPLOAD_ACCESS_KEY_ID - value: minioadmin + value: MINIO_USER_PLACEHOLDER - name: LANGFUSE_S3_EVENT_UPLOAD_SECRET_ACCESS_KEY - value: minioadmin + value: MINIO_PASSWORD_PLACEHOLDER - name: LANGFUSE_S3_EVENT_UPLOAD_ENDPOINT value: http://127.0.0.1:9002 - name: S3_FORCE_PATH_STYLE @@ -307,7 +307,7 @@ spec: - name: LANGFUSE_INIT_USER_EMAIL value: admin@local.dev - name: LANGFUSE_INIT_USER_PASSWORD - value: admin-local-pw-2026 + value: LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER - name: LANGFUSE_LOG_LEVEL value: warn image: docker.io/langfuse/langfuse:3 @@ -363,13 +363,13 @@ spec: - name: LANGFUSE_S3_EVENT_UPLOAD_REGION value: us-east-1 - name: LANGFUSE_S3_EVENT_UPLOAD_ACCESS_KEY_ID - value: minioadmin + value: MINIO_USER_PLACEHOLDER - name: S3_FORCE_PATH_STYLE value: 'true' - name: LANGFUSE_S3_EVENT_UPLOAD_ENDPOINT value: http://127.0.0.1:9002 - name: LANGFUSE_S3_EVENT_UPLOAD_SECRET_ACCESS_KEY - value: minioadmin + value: MINIO_PASSWORD_PLACEHOLDER - name: LANGFUSE_LOG_LEVEL value: warn image: docker.io/langfuse/langfuse-worker:3 @@ -393,9 +393,9 @@ spec: - ":9001" env: - name: MINIO_ROOT_USER - value: minioadmin + value: MINIO_USER_PLACEHOLDER - name: MINIO_ROOT_PASSWORD - value: minioadmin + value: MINIO_PASSWORD_PLACEHOLDER image: docker.io/minio/minio:latest livenessProbe: httpGet: diff --git a/router/Dockerfile b/router/Dockerfile index cd0f3653..49d9c6a6 100644 --- a/router/Dockerfile +++ b/router/Dockerfile @@ -3,7 +3,7 @@ FROM python:3.14-slim WORKDIR /app # Install deps in a single layer (no pip cache, no extra files) -RUN pip install --no-cache-dir fastapi "pydantic>=2.0,<3.0" uvicorn httpx pyyaml python-multipart asyncpg langfuse redis +RUN pip install --no-cache-dir fastapi "pydantic>=2.0,<3.0" uvicorn httpx pyyaml python-multipart asyncpg langfuse redis aiofiles # Copy all source in one layer — removes dead config COPY (volume-mounted at runtime) COPY main.py agy_proxy.py circuit_breaker.py aa_scores.json free_models_roster.json /app/ diff --git a/router/main.py b/router/main.py index 6c1c9a2d..9e379255 100644 --- a/router/main.py +++ b/router/main.py @@ -1,4 +1,6 @@ import os +import aiofiles +import re import sys import json import time @@ -49,10 +51,15 @@ def get_redis(): return None _redis_last_init_attempt = now try: - host = os.getenv("VALKEY_HOST", "127.0.0.1") - port = int(os.getenv("VALKEY_PORT", "6379")) - _redis_client = aioredis.Redis(host=host, port=port, decode_responses=True, socket_timeout=1.0) - logger.info(f"Valkey client initialized at {host}:{port}") + url = os.getenv("VALKEY_URL") + if url: + _redis_client = aioredis.Redis.from_url(url, decode_responses=True, socket_timeout=1.0) + logger.info("Valkey client initialized from URL") + else: + host = os.getenv("VALKEY_HOST", "127.0.0.1") + port = int(os.getenv("VALKEY_PORT", "6379")) + _redis_client = aioredis.Redis(host=host, port=port, decode_responses=True, socket_timeout=1.0) + logger.info(f"Valkey client initialized at {host}:{port}") except Exception as e: logger.warning(f"Failed to initialize Valkey client: {e} — falling back to local memory") _redis_client = None @@ -82,24 +89,60 @@ def get_http_client(): return _http_client +# Compiled regular expressions for token estimation heuristics +WORD_RE = re.compile(r'[a-zA-Z0-9]+') +NON_ASCII_RE = re.compile(r'[^\s\x00-\x7F]') +PUNC_RE = re.compile(r'[\x21-\x2f\x3a-\x40\x5b-\x60\x7b-\x7e]') + + +def _count_tokens_heuristic(text: str) -> float: + """Heuristically estimate token count using weighted categories and optimized regex splitting. + + This replaces the naive character-count logic with a more granular approach that + balances English words, technical identifiers, punctuation, and multi-byte characters. + + Returns a float to prevent intermediate rounding errors when summing across multiple + message blocks. Callers should round the total sum to convert it to an integer. + """ + if not text: + return 0.0 + + # 1. Alphanumeric runs (Words/Identifiers/Hashes/Base64) + # Use a length-aware heuristic to avoid under-counting technical content. + word_matches = WORD_RE.findall(text) + word_total = sum(1.2 if len(w) <= 8 else len(w) / 4.0 for w in word_matches) + + # 2. Non-ASCII characters (CJK/Emoji) + # Each character is weighted at 0.35 tokens. + non_ascii_count = len(NON_ASCII_RE.findall(text)) + + # 3. ASCII Punctuation/Symbols + # Characters that are ASCII but not alphanumeric or whitespace. + punc_count = len(PUNC_RE.findall(text)) + + return word_total + (non_ascii_count * 0.35) + (punc_count * 0.4) + + def estimate_prompt_tokens(body: dict) -> int: - """Estimate prompt tokens by counting characters in message contents (1 token ~= 4 chars) - to avoid inflating metrics with large tool/schema declarations. + """Estimate prompt tokens using a regex-based weighted heuristic for mixed content. """ - tokens = 0 + total = 0.0 for msg in body.get("messages", []): if not isinstance(msg, dict): continue content = msg.get("content") or "" if isinstance(content, str): - tokens += len(content) // 4 + total += _count_tokens_heuristic(content) elif isinstance(content, list): for block in content: if isinstance(block, dict) and block.get("type") == "text": - tokens += len(block.get("text") or "") // 4 - # Include a flat estimate for system prompt / metadata overhead - tokens += 50 - return max(1, tokens) + text = block.get("text") + if isinstance(text, str): + total += _count_tokens_heuristic(text) + + # Include a flat estimate for system prompt / metadata overhead. + # Use rounding to avoid truncation bias (e.g., 1.9 -> 1). + return max(1, int(round(total)) + 50) async def sync_cooldowns_from_valkey() -> None: @@ -3871,9 +3914,26 @@ def validate_payload(self) -> "AnnotationPayload": annotations_lock = asyncio.Lock() -def _read_annotations_sync(path) -> dict: - with open(path, "r", encoding="utf-8") as f: - return json.load(f) +_annotations_cache = {} + + +async def _read_annotations_async(path) -> dict: + import copy + + # Do not swallow OSError if file doesn't exist to preserve original behavior. + # The caller (save_annotations) handles the exception when reading existing annotations. + current_mtime = await asyncio.to_thread(os.path.getmtime, path) + + cache_entry = _annotations_cache.get(path) + + if cache_entry is None or current_mtime != cache_entry["mtime"]: + async with aiofiles.open(path, "r", encoding="utf-8") as f: + # Read asynchronously, but parse in a thread pool to avoid blocking event loop + content = await f.read() + data = await asyncio.to_thread(json.loads, content) + _annotations_cache[path] = {"mtime": current_mtime, "data": data} + + return await asyncio.to_thread(copy.deepcopy, _annotations_cache[path]["data"]) @app.post("/dashboard/save-annotations") @@ -3887,9 +3947,7 @@ async def save_annotations(payload: AnnotationPayload): async with annotations_lock: if ann_path.exists(): try: - existing = await asyncio.to_thread( - _read_annotations_sync, str(ann_path) - ) + existing = await _read_annotations_async(str(ann_path)) except Exception as read_err: logger.warning( f"Could not read existing annotations: {read_err}. Overwriting." diff --git a/router/memory_mcp.py b/router/memory_mcp.py index f2187282..7ec87c26 100755 --- a/router/memory_mcp.py +++ b/router/memory_mcp.py @@ -17,6 +17,7 @@ import time import hashlib import httpx +import urllib.parse API_URL = "http://127.0.0.1:5000/v1/memory" PROTOCOL_VERSION = "2024-11-05" @@ -44,7 +45,8 @@ def _make_key(category: str, is_global: bool, data: str) -> str: # BLAKE2b: SOTA crypto hash, stdlib, faster than MD5, deterministic across restarts. # Provides uniqueness within the same millisecond. h = hashlib.blake2b((data + str(ts)).encode("utf-8"), digest_size=HASH_DIGEST_SIZE).hexdigest() - return f"{PREFIX}:{scope}:{category}::{ts}:{h}" + safe_category = urllib.parse.quote(category) + return f"{PREFIX}:{scope}:{safe_category}::{ts}:{h}" def _parse_key(key: str): @@ -53,7 +55,7 @@ def _parse_key(key: str): parts = key.split("::") prefix = parts[0].split(":") # memory:{scope}:{category} scope = prefix[1] if len(prefix) > 1 else "" - category = prefix[2] if len(prefix) > 2 else "" + category = urllib.parse.unquote(prefix[2]) if len(prefix) > 2 else "" ts_hash = parts[1] if len(parts) > 1 else "" ts = ts_hash.split(":")[0] if ts_hash else "" return {"scope": scope, "category": category, "timestamp": ts} @@ -68,6 +70,8 @@ def _parse_key(key: str): def _is_memory_key(key: str) -> bool: """Check if a key follows the memory:{scope}:{category}:: format.""" + if not isinstance(key, str): + return False return key.startswith(f"{PREFIX}:") diff --git a/router/test_memory_mcp.py b/router/test_memory_mcp.py deleted file mode 100644 index 24d23a34..00000000 --- a/router/test_memory_mcp.py +++ /dev/null @@ -1,131 +0,0 @@ -import time -import re -import sys -import json -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from memory_mcp import _make_key, SCOPE_GLOBAL, SCOPE_LOCAL, PREFIX, _memory_value, _parse_memory_value - -def test_make_key_global(): - """Test generating a key for global scope.""" - category = "test_cat" - data = "test_data" - - before_ts = int(time.time() * 1000) - key = _make_key(category, True, data) - after_ts = int(time.time() * 1000) - - # Expected format: f"{PREFIX}:{scope}:{category}::{ts}:{h}" - parts = key.split(":") - - assert key.startswith(f"{PREFIX}:{SCOPE_GLOBAL}:{category}::") - - # Extract timestamp and hash part - # Format is memory:global:test_cat::1717612345:a1b2c3d4... - match = re.match(rf"^{PREFIX}:{SCOPE_GLOBAL}:{category}::(\d+):([a-f0-9]+)$", key) - assert match is not None, f"Key {key} does not match expected format" - - ts = int(match.group(1)) - h = match.group(2) - - assert before_ts <= ts <= after_ts - assert len(h) == 20 - -def test_make_key_local(): - """Test generating a key for local scope.""" - category = "another_cat" - data = "more_data" - - before_ts = int(time.time() * 1000) - key = _make_key(category, False, data) - after_ts = int(time.time() * 1000) - - assert key.startswith(f"{PREFIX}:{SCOPE_LOCAL}:{category}::") - - match = re.match(rf"^{PREFIX}:{SCOPE_LOCAL}:{category}::(\d+):([a-f0-9]+)$", key) - assert match is not None, f"Key {key} does not match expected format" - - ts = int(match.group(1)) - h = match.group(2) - - assert before_ts <= ts <= after_ts - assert len(h) == 20 - - -def test_make_key_formatting_details(monkeypatch): - """Test the exact output formatting of _make_key using deterministic BLAKE2b.""" - # Mock time.time to return a predictable float so ts = 1620000000123 - monkeypatch.setattr(time, "time", lambda: 1620000000.123) - - # data="data", ts=1620000000123 -> blake2b("data1620000000123", digest_size=10) -> 5e5dad075ca7764bc51f - key1 = _make_key("cat1", True, "data") - assert key1 == f"{PREFIX}:{SCOPE_GLOBAL}:cat1::1620000000123:5e5dad075ca7764bc51f" - - key2 = _make_key("cat2", False, "data") - assert key2 == f"{PREFIX}:{SCOPE_LOCAL}:cat2::1620000000123:5e5dad075ca7764bc51f" - - -def test_make_key_determinism_and_uniqueness(): - """Test determinism for same inputs within same timestamp, and uniqueness across timestamps/data.""" - category = "test_cat" - data1 = "data1" - data2 = "data2" - - key1 = _make_key(category, True, data1) - time.sleep(0.002) - key2 = _make_key(category, True, data1) - key3 = _make_key(category, True, data2) - - # Uniqueness across data - assert key1 != key3 - - # Check determinism: if the timestamp parts are the same, the keys should be identical - ts1 = key1.split("::")[1].split(":")[0] - ts2 = key2.split("::")[1].split(":")[0] - if ts1 == ts2: - assert key1 == key2 - else: - # If timestamp is different, keys should be different - assert key1 != key2 - -def test_memory_value_happy_path(): - """Test _memory_value with standard data and tags.""" - result = _memory_value("some data", ["tag1", "tag2"]) - parsed = json.loads(result) - assert parsed == {"data": "some data", "tags": ["tag1", "tag2"]} - -def test_memory_value_missing_tags(): - """Test _memory_value when tags is None.""" - result = _memory_value("some data", None) - parsed = json.loads(result) - assert parsed == {"data": "some data", "tags": []} - -def test_memory_value_unicode(): - """Test _memory_value properly handles unicode and ensure_ascii=False.""" - result = _memory_value("こんにちは", ["世界"]) - # If ensure_ascii=False, the unicode characters shouldn't be escaped (no \uXXXX) - assert "こんにちは" in result - assert "世界" in result - parsed = json.loads(result) - assert parsed == {"data": "こんにちは", "tags": ["世界"]} - -def test_parse_memory_value_success(): - """Test _parse_memory_value successfully decodes valid JSON.""" - raw = '{"data": "info", "tags": ["a"]}' - result = _parse_memory_value(raw) - assert result == {"data": "info", "tags": ["a"]} - -def test_parse_memory_value_invalid_json(): - """Test _parse_memory_value with invalid JSON.""" - result = _parse_memory_value("{invalid_json:") - assert result == {"data": "{invalid_json:", "tags": []} - -def test_parse_memory_value_type_error(): - """Test _parse_memory_value with TypeError (e.g. passing None).""" - result = _parse_memory_value(None) - assert result == {"data": None, "tags": []} - -def test_parse_memory_value_invalid_json_string(): - """Test _parse_memory_value with invalid JSON string.""" - result = _parse_memory_value("this is not a valid json string") - assert result == {"data": "this is not a valid json string", "tags": []} diff --git a/router/tests/test_detect_active_tool.py b/router/tests/test_detect_active_tool.py new file mode 100644 index 00000000..3105ab1e --- /dev/null +++ b/router/tests/test_detect_active_tool.py @@ -0,0 +1,114 @@ +import pytest +import os +import sys +from pathlib import Path + +# Set CONFIG_PATH for import +os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "config.yaml") + +# Add the parent directory to the path so we can import from router +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from router.main import detect_active_tool + +def test_detect_active_tool_empty(): + assert detect_active_tool({}) == "none" + assert detect_active_tool({"messages": []}) == "none" + assert detect_active_tool({"messages": [{"role": "system", "content": "hello"}]}) == "none" + +def test_detect_active_tool_role_tool_with_name(): + # Write tool name mapped to write + body = { + "messages": [ + {"role": "tool", "name": "edit_file", "content": "..."} + ] + } + assert detect_active_tool(body) == "write" + + # View tool mapped to view + body = { + "messages": [ + {"role": "tool", "name": "cat_file", "content": "..."} + ] + } + assert detect_active_tool(body) == "view" + +def test_detect_active_tool_role_tool_without_name_but_matched_tool_call_id(): + body = { + "messages": [ + {"role": "assistant", "tool_calls": [{"id": "call_123", "function": {"name": "read_file"}}]}, + {"role": "tool", "tool_call_id": "call_123", "content": "success"} + ] + } + assert detect_active_tool(body) == "view" + +def test_detect_active_tool_role_tool_without_name_unmatched_tool_call_id(): + body = { + "messages": [ + {"role": "assistant", "tool_calls": [{"id": "call_999", "function": {"name": "read_file"}}]}, + {"role": "tool", "tool_call_id": "call_123", "content": "success"} + ] + } + assert detect_active_tool(body) == "other" + +def test_detect_active_tool_role_assistant_with_tool_calls(): + body = { + "messages": [ + {"role": "user", "content": "do something"}, + {"role": "assistant", "tool_calls": [{"id": "call_1", "function": {"name": "write_to_file"}}]} + ] + } + assert detect_active_tool(body) == "write" + +def test_detect_active_tool_fallback_user_keyword(): + # Tests matching "tree" + body = { + "messages": [ + {"role": "user", "content": "show me the tree"} + ] + } + assert detect_active_tool(body) == "tree" + + # Tests matching "shell" + body = { + "messages": [ + {"role": "user", "content": "run this in shell"} + ] + } + assert detect_active_tool(body) == "shell" + + # Tests matching "write" + body = { + "messages": [ + {"role": "user", "content": "create file test.py"} + ] + } + assert detect_active_tool(body) == "write" + + # Tests matching "view" + body = { + "messages": [ + {"role": "user", "content": "cat main.py"} + ] + } + assert detect_active_tool(body) == "view" + +def test_detect_active_tool_ignores_invalid_message_formats(): + body = { + "messages": [ + "this is not a dict", + {"role": "user", "content": "read this"} + ] + } + assert detect_active_tool(body) == "view" + +def test_detect_active_tool_precedence(): + # If there are multiple tools, it processes from the last message backwards + body = { + "messages": [ + {"role": "tool", "name": "edit_file", "content": "done"}, + {"role": "tool", "name": "cat_file", "content": "..."} + ] + } + # It starts from the end, so it sees "cat_file" first, which maps to "view" + assert detect_active_tool(body) == "view" diff --git a/router/tests/test_estimate_prompt_tokens.py b/router/tests/test_estimate_prompt_tokens.py index e93390f9..ae74a9e5 100644 --- a/router/tests/test_estimate_prompt_tokens.py +++ b/router/tests/test_estimate_prompt_tokens.py @@ -23,25 +23,27 @@ def test_estimate_prompt_tokens_empty_messages(): def test_estimate_prompt_tokens_string_content(): body = { "messages": [ - {"content": "1234"}, # 1 token - {"content": "12345678"} # 2 tokens + {"content": "word " * 4}, # 4 * 1.2 = 4.8 + {"content": "word " * 8} # 8 * 1.2 = 9.6 ] } - assert estimate_prompt_tokens(body) == 50 + 1 + 2 + # Total is int(round(4.8 + 9.6)) + 50 = int(round(14.4)) + 50 = 14 + 50 = 64 + assert estimate_prompt_tokens(body) == 50 + 14 def test_estimate_prompt_tokens_list_content(): body = { "messages": [ { "content": [ - {"type": "text", "text": "1234"}, # 1 token + {"type": "text", "text": "word " * 4}, # 4.8 tokens {"type": "image_url", "url": "ignored"}, # 0 tokens - {"type": "text", "text": "12345678"} # 2 tokens + {"type": "text", "text": "word " * 8} # 9.6 tokens ] } ] } - assert estimate_prompt_tokens(body) == 50 + 1 + 2 + # Total is int(round(4.8 + 9.6)) + 50 = 64 + assert estimate_prompt_tokens(body) == 50 + 14 def test_estimate_prompt_tokens_mixed_and_invalid_msgs(): body = { @@ -52,10 +54,11 @@ def test_estimate_prompt_tokens_mixed_and_invalid_msgs(): "invalid_block_type", # Should be skipped {"type": "text", "text": None} # None text, handled as empty string ]}, - {"content": "1234"} # 1 token + {"content": "word " * 4} # 4.8 tokens ] } - assert estimate_prompt_tokens(body) == 50 + 1 + # Total is int(round(4.8)) + 50 = 5 + 50 = 55 + assert estimate_prompt_tokens(body) == 50 + 5 def test_estimate_prompt_tokens_missing_content(): body = { diff --git a/router/tests/test_get_gemini_oauth_status.py b/router/tests/test_get_gemini_oauth_status.py new file mode 100644 index 00000000..89523237 --- /dev/null +++ b/router/tests/test_get_gemini_oauth_status.py @@ -0,0 +1,99 @@ +import pytest +import os +import sys +import json +from unittest.mock import patch, mock_open + +# Ensure router directory is in sys.path +router_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if router_path not in sys.path: + sys.path.insert(0, router_path) + +import main + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_missing_file(): + with patch("os.path.exists", return_value=False): + result = await main.get_gemini_oauth_status() + assert result == {"status": "missing", "detail": "No oauth_creds.json found", "expiry_ms": 0} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_no_access_token(): + mock_data = {"expiry_date": 1234567890000} + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", mock_open(read_data=json.dumps(mock_data))): + result = await main.get_gemini_oauth_status() + assert result == {"status": "missing", "detail": "No access token in file", "expiry_ms": 0} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_valid_less_than_60s(): + current_time_ms = 1000000000000 + expiry_ms = current_time_ms + 45000 # 45 seconds from now + mock_data = {"access_token": "test_token", "expiry_date": expiry_ms} + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", mock_open(read_data=json.dumps(mock_data))), \ + patch("time.time", return_value=current_time_ms / 1000.0): + result = await main.get_gemini_oauth_status() + assert result == {"status": "valid", "detail": "Expires in 45s", "expiry_ms": expiry_ms} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_valid_less_than_3600s(): + current_time_ms = 1000000000000 + expiry_ms = current_time_ms + 1500000 # 25 minutes from now (25 * 60 = 1500 seconds) + mock_data = {"access_token": "test_token", "expiry_date": expiry_ms} + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", mock_open(read_data=json.dumps(mock_data))), \ + patch("time.time", return_value=current_time_ms / 1000.0): + result = await main.get_gemini_oauth_status() + assert result == {"status": "valid", "detail": "Expires in 25m 0s", "expiry_ms": expiry_ms} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_valid_more_than_3600s(): + current_time_ms = 1000000000000 + expiry_ms = current_time_ms + 7500000 # 2 hours, 5 minutes from now (2 * 3600 + 5 * 60 = 7500) + mock_data = {"access_token": "test_token", "expiry_date": expiry_ms} + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", mock_open(read_data=json.dumps(mock_data))), \ + patch("time.time", return_value=current_time_ms / 1000.0): + result = await main.get_gemini_oauth_status() + assert result == {"status": "valid", "detail": "Expires in 2h 5m", "expiry_ms": expiry_ms} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_expired_less_than_3600s(): + current_time_ms = 1000000000000 + expiry_ms = current_time_ms - 1500000 # Expired 25 minutes ago + mock_data = {"access_token": "test_token", "expiry_date": expiry_ms} + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", mock_open(read_data=json.dumps(mock_data))), \ + patch("time.time", return_value=current_time_ms / 1000.0): + result = await main.get_gemini_oauth_status() + assert result == {"status": "expired", "detail": "Expired 25 minutes ago", "expiry_ms": expiry_ms} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_expired_less_than_86400s(): + current_time_ms = 1000000000000 + expiry_ms = current_time_ms - 7500000 # Expired 2 hours ago + mock_data = {"access_token": "test_token", "expiry_date": expiry_ms} + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", mock_open(read_data=json.dumps(mock_data))), \ + patch("time.time", return_value=current_time_ms / 1000.0): + result = await main.get_gemini_oauth_status() + assert result == {"status": "expired", "detail": "Expired 2 hours ago", "expiry_ms": expiry_ms} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_expired_more_than_86400s(): + current_time_ms = 1000000000000 + expiry_ms = current_time_ms - 172800000 # Expired 2 days ago (2 * 86400 * 1000 = 172800000) + mock_data = {"access_token": "test_token", "expiry_date": expiry_ms} + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", mock_open(read_data=json.dumps(mock_data))), \ + patch("time.time", return_value=current_time_ms / 1000.0): + result = await main.get_gemini_oauth_status() + assert result == {"status": "expired", "detail": "Expired 2 days ago", "expiry_ms": expiry_ms} + +@pytest.mark.asyncio +async def test_get_gemini_oauth_status_exception(): + with patch("os.path.exists", return_value=True), \ + patch("builtins.open", side_effect=Exception("Test error")): + result = await main.get_gemini_oauth_status() + assert result == {"status": "error", "detail": "Test error", "expiry_ms": 0} diff --git a/router/tests/test_get_goose_sessions.py b/router/tests/test_get_goose_sessions.py new file mode 100644 index 00000000..52640fef --- /dev/null +++ b/router/tests/test_get_goose_sessions.py @@ -0,0 +1,46 @@ +import pytest +import sys +import os +from pathlib import Path +from unittest.mock import patch, MagicMock + +os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "config.yaml") +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from main import get_goose_sessions + +def test_get_goose_sessions_no_db(): + with patch('os.path.exists', return_value=False): + assert get_goose_sessions() == [] + +def test_get_goose_sessions_success(): + mock_sqlite3 = MagicMock() + mock_conn = MagicMock() + mock_cursor = MagicMock() + + mock_sqlite3.connect.return_value = mock_conn + mock_conn.cursor.return_value = mock_cursor + + mock_cursor.fetchall.return_value = [ + {"id": 1, "name": "s1"}, + {"id": 2, "name": "s2"} + ] + + with patch('os.path.exists', return_value=True): + with patch.dict(sys.modules, {'sqlite3': mock_sqlite3}): + result = get_goose_sessions() + + assert len(result) == 2 + assert result[0] == {"id": 1, "name": "s1"} + mock_sqlite3.connect.assert_called_once_with("/config/goose_sessions/sessions/sessions.db", timeout=1.0) + mock_cursor.execute.assert_called_once() + mock_conn.close.assert_called_once() + +def test_get_goose_sessions_exception(): + mock_sqlite3 = MagicMock() + mock_sqlite3.connect.side_effect = Exception("DB error") + + with patch('os.path.exists', return_value=True): + with patch.dict(sys.modules, {'sqlite3': mock_sqlite3}): + result = get_goose_sessions() + assert result == [] diff --git a/router/tests/test_get_redis.py b/router/tests/test_get_redis.py new file mode 100644 index 00000000..0fdeaed2 --- /dev/null +++ b/router/tests/test_get_redis.py @@ -0,0 +1,136 @@ +import os +import time +from unittest.mock import patch, MagicMock + +import pytest + +import router.main as main + +@pytest.fixture(autouse=True) +def reset_redis_globals(): + """Reset the global variables before and after each test.""" + original_client = main._redis_client + original_last_attempt = main._redis_last_init_attempt + + main._redis_client = None + main._redis_last_init_attempt = 0.0 + + yield + + main._redis_client = original_client + main._redis_last_init_attempt = original_last_attempt + +def test_get_redis_already_initialized(): + """If the client is already initialized, it should return the client immediately.""" + mock_client = MagicMock() + main._redis_client = mock_client + + assert main.get_redis() is mock_client + +@patch("router.main.time.monotonic") +def test_get_redis_cooldown(mock_monotonic): + """If init failed recently, it should return None without attempting to initialize.""" + main._redis_client = None + main._redis_last_init_attempt = 100.0 + + # Time elapsed is less than 5.0 seconds + mock_monotonic.return_value = 103.0 + + assert main.get_redis() is None + +@patch("router.main.time.monotonic") +@patch("router.main.aioredis.Redis") +@patch.dict(os.environ, {"VALKEY_HOST": "my-host", "VALKEY_PORT": "1234"}) +def test_get_redis_initialization_success(mock_redis, mock_monotonic): + """If sufficient time has passed, it should initialize and return the client.""" + main._redis_client = None + main._redis_last_init_attempt = 100.0 + + # Time elapsed is 10.0 seconds (greater than 5.0) + mock_monotonic.return_value = 110.0 + + mock_redis_instance = MagicMock() + mock_redis.return_value = mock_redis_instance + + client = main.get_redis() + + assert client is mock_redis_instance + assert main._redis_client is mock_redis_instance + assert main._redis_last_init_attempt == 110.0 + mock_redis.assert_called_once_with(host="my-host", port=1234, decode_responses=True, socket_timeout=1.0) + +@patch("router.main.time.monotonic") +@patch("router.main.logger.warning") +@patch.dict(os.environ, {"VALKEY_HOST": "my-host", "VALKEY_PORT": "invalid"}) +def test_get_redis_initialization_failure(mock_logger_warning, mock_monotonic): + """If initialization fails, it should catch the exception, log a warning, and return None.""" + main._redis_client = None + main._redis_last_init_attempt = 100.0 + + # Time elapsed is 10.0 seconds + mock_monotonic.return_value = 110.0 + + # The int() conversion of VALKEY_PORT will raise ValueError + client = main.get_redis() + + assert client is None + assert main._redis_client is None + assert main._redis_last_init_attempt == 110.0 + mock_logger_warning.assert_called_once() + assert "Failed to initialize Valkey client" in mock_logger_warning.call_args[0][0] + +@patch("router.main.time.monotonic") +@patch("router.main.aioredis.Redis") +@patch("router.main.logger.warning") +@patch.dict(os.environ, {"VALKEY_HOST": "my-host", "VALKEY_PORT": "1234"}) +def test_get_redis_initialization_exception(mock_logger_warning, mock_redis, mock_monotonic): + """If aioredis.Redis throws an exception, it should catch it and return None.""" + main._redis_client = None + main._redis_last_init_attempt = 100.0 + + mock_monotonic.return_value = 110.0 + + mock_redis.side_effect = Exception("Test Exception") + + client = main.get_redis() + + assert client is None + assert main._redis_client is None + assert main._redis_last_init_attempt == 110.0 + mock_logger_warning.assert_called_once() + assert "Test Exception" in mock_logger_warning.call_args[0][0] + +@patch("router.main.time.monotonic") +@patch("router.main.aioredis.Redis.from_url") +@patch.dict(os.environ, {"VALKEY_URL": "redis://my-url:1234"}) +def test_get_redis_simulation_flow_url(mock_from_url, mock_monotonic): + """Simulate the full flow for from_url: failure -> cooldown -> success -> cached.""" + # State is reset by the autouse fixture reset_redis_globals + + # 1. First attempt fails + mock_monotonic.return_value = 10.0 + mock_from_url.side_effect = Exception("Connection error") + assert main.get_redis() is None + assert main._redis_last_init_attempt == 10.0 + + # 2. Second attempt during cooldown (e.g. 12.0s) + mock_monotonic.return_value = 12.0 + mock_from_url.reset_mock() + assert main.get_redis() is None + mock_from_url.assert_not_called() + + # 3. Third attempt after cooldown (e.g. 16.0s) succeeds + mock_monotonic.return_value = 16.0 + mock_redis_instance = MagicMock() + mock_from_url.side_effect = None + mock_from_url.return_value = mock_redis_instance + client = main.get_redis() + assert client is mock_redis_instance + assert main._redis_client is mock_redis_instance + mock_from_url.assert_called_once() + + # 4. Fourth attempt returns cached instance + mock_monotonic.return_value = 18.0 + mock_from_url.reset_mock() + assert main.get_redis() is mock_redis_instance + mock_from_url.assert_not_called() diff --git a/router/tests/test_load_persisted_stats.py b/router/tests/test_load_persisted_stats.py new file mode 100644 index 00000000..76edbf51 --- /dev/null +++ b/router/tests/test_load_persisted_stats.py @@ -0,0 +1,61 @@ +import json +import pytest +from unittest.mock import patch, mock_open + +import router.main +from router.main import load_persisted_stats + +@pytest.fixture +def mock_stats(): + # Setup a clean stats dictionary for testing + clean_stats = { + "total_requests": 0, + "nested_dict": {"a": 1, "b": 2}, + "existing_key": "value" + } + with patch.dict(router.main.stats, clean_stats, clear=True): + yield router.main.stats + +def test_load_persisted_stats_file_not_exists(mock_stats): + with patch("router.main.os.path.exists", return_value=False) as mock_exists: + load_persisted_stats() + mock_exists.assert_called_once_with(router.main.STATS_JSON_PATH) + # Stats should remain unchanged + assert mock_stats["total_requests"] == 0 + +def test_load_persisted_stats_success(mock_stats): + mock_data = { + "total_requests": 100, + "nested_dict": {"b": 3, "c": 4}, + "new_key": "new_value" + } + mock_json = json.dumps(mock_data) + + with patch("router.main.os.path.exists", return_value=True): + with patch("router.main.open", mock_open(read_data=mock_json)): + with patch("router.main.logger.info") as mock_logger: + load_persisted_stats() + + # Assert simple value updated via else block + assert mock_stats["total_requests"] == 100 + # Assert nested_dict updated via if block (b updated, c added, a unchanged) + assert mock_stats["nested_dict"] == {"a": 1, "b": 3, "c": 4} + # Assert new_key added via else block + assert mock_stats["new_key"] == "new_value" + # Assert existing_key unchanged + assert mock_stats["existing_key"] == "value" + + mock_logger.assert_called_once_with("✓ Successfully loaded persisted gateway statistics from disk.") + +def test_load_persisted_stats_exception(mock_stats): + with patch("router.main.os.path.exists", return_value=True): + with patch("router.main.open", side_effect=Exception("Mock read error")): + with patch("router.main.logger.error") as mock_logger: + load_persisted_stats() + + # Stats should remain unchanged + assert mock_stats["total_requests"] == 0 + + # Error should be logged + mock_logger.assert_called_once() + assert "Failed to load persisted stats: Mock read error" in mock_logger.call_args[0][0] diff --git a/router/tests/test_memory_mcp.py b/router/tests/test_memory_mcp.py new file mode 100644 index 00000000..8c93f1f6 --- /dev/null +++ b/router/tests/test_memory_mcp.py @@ -0,0 +1,313 @@ +import json +import re +import sys +import time +from pathlib import Path + +# Dynamic project root discovery +root = Path(__file__).resolve() +while root.parent != root and not (root / ".git").exists(): + root = root.parent +sys.path.insert(0, str(root)) +sys.path.insert(0, str(root / "router")) + +import pytest +from memory_mcp import ( + PREFIX, + SCOPE_GLOBAL, + SCOPE_LOCAL, + _make_key, + _memory_entry, + _memory_value, + _parse_key, + _parse_memory_value, +) + + +# ===================================================================== +# Tests from router/test_memory_mcp.py +# ===================================================================== + +def test_make_key_global(): + """Test generating a key for global scope.""" + category = "test_cat" + data = "test_data" + + before_ts = int(time.time() * 1000) + key = _make_key(category, True, data) + after_ts = int(time.time() * 1000) + + # Expected format: f"{PREFIX}:{scope}:{category}::{ts}:{h}" + assert key.startswith(f"{PREFIX}:{SCOPE_GLOBAL}:{category}::") + + # Extract timestamp and hash part + match = re.match(rf"^{PREFIX}:{SCOPE_GLOBAL}:{category}::(\d+):([a-f0-9]+)$", key) + assert match is not None, f"Key {key} does not match expected format" + + ts = int(match.group(1)) + h = match.group(2) + + assert before_ts <= ts <= after_ts + assert len(h) == 20 + + +def test_make_key_local(): + """Test generating a key for local scope.""" + category = "another_cat" + data = "more_data" + + before_ts = int(time.time() * 1000) + key = _make_key(category, False, data) + after_ts = int(time.time() * 1000) + + assert key.startswith(f"{PREFIX}:{SCOPE_LOCAL}:{category}::") + + match = re.match(rf"^{PREFIX}:{SCOPE_LOCAL}:{category}::(\d+):([a-f0-9]+)$", key) + assert match is not None, f"Key {key} does not match expected format" + + ts = int(match.group(1)) + h = match.group(2) + + assert before_ts <= ts <= after_ts + assert len(h) == 20 + + +def test_make_key_formatting_details(monkeypatch): + """Test the exact output formatting of _make_key using deterministic BLAKE2b.""" + # Mock time.time to return a predictable float so ts = 1620000000123 + monkeypatch.setattr(time, "time", lambda: 1620000000.123) + + # data="data", ts=1620000000123 -> blake2b("data1620000000123", digest_size=10) -> 5e5dad075ca7764bc51f + key1 = _make_key("cat1", True, "data") + assert key1 == f"{PREFIX}:{SCOPE_GLOBAL}:cat1::1620000000123:5e5dad075ca7764bc51f" + + key2 = _make_key("cat2", False, "data") + assert key2 == f"{PREFIX}:{SCOPE_LOCAL}:cat2::1620000000123:5e5dad075ca7764bc51f" + + +def test_make_key_determinism_and_uniqueness(): + """Test determinism for same inputs within same timestamp, and uniqueness across timestamps/data.""" + category = "test_cat" + data1 = "data1" + data2 = "data2" + + key1 = _make_key(category, True, data1) + time.sleep(0.002) + key2 = _make_key(category, True, data1) + key3 = _make_key(category, True, data2) + + # Uniqueness across data + assert key1 != key3 + + # Check determinism: if the timestamp parts are the same, the keys should be identical + ts1 = key1.split("::")[1].split(":")[0] + ts2 = key2.split("::")[1].split(":")[0] + if ts1 == ts2: + assert key1 == key2 + else: + # If timestamp is different, keys should be different + assert key1 != key2 + + +def test_memory_value_happy_path(): + """Test _memory_value with standard data and tags.""" + result = _memory_value("some data", ["tag1", "tag2"]) + parsed = json.loads(result) + assert parsed == {"data": "some data", "tags": ["tag1", "tag2"]} + + +def test_memory_value_missing_tags(): + """Test _memory_value when tags is None.""" + result = _memory_value("some data", None) + parsed = json.loads(result) + assert parsed == {"data": "some data", "tags": []} + + +def test_memory_value_unicode(): + """Test _memory_value properly handles unicode and ensure_ascii=False.""" + result = _memory_value("こんにちは", ["世界"]) + # If ensure_ascii=False, the unicode characters shouldn't be escaped (no \uXXXX) + assert "こんにちは" in result + assert "世界" in result + parsed = json.loads(result) + assert parsed == {"data": "こんにちは", "tags": ["世界"]} + + +def test_parse_memory_value_success(): + """Test _parse_memory_value successfully decodes valid JSON.""" + raw = '{"data": "info", "tags": ["a"]}' + result = _parse_memory_value(raw) + assert result == {"data": "info", "tags": ["a"]} + + +def test_parse_memory_value_invalid_json(): + """Test _parse_memory_value with invalid JSON.""" + result = _parse_memory_value("{invalid_json:") + assert result == {"data": "{invalid_json:", "tags": []} + + +def test_parse_memory_value_type_error(): + """Test _parse_memory_value with TypeError (e.g. passing None).""" + result = _parse_memory_value(None) + assert result == {"data": None, "tags": []} + + +def test_parse_memory_value_invalid_json_string(): + """Test _parse_memory_value with invalid JSON string.""" + result = _parse_memory_value("this is not a valid json string") + assert result == {"data": "this is not a valid json string", "tags": []} + + +# ===================================================================== +# Tests from test_memory_mcp.py (root) +# ===================================================================== + +def test_memory_entry_happy_path(): + """Test correctly formatted and complete memory entry.""" + valid_key = "memory:global:project_standards::1689201948123:a1b2c3d4e5f6" + valid_value = json.dumps({"data": "Use pytest for all tests", "tags": ["testing", "python"]}) + lmem = { + "key": valid_key, + "value": valid_value, + "memory_id": "test_id_123" + } + + result = _memory_entry(lmem) + + assert result is not None + assert result["key"] == valid_key + assert result["category"] == "project_standards" + assert result["data"] == "Use pytest for all tests" + assert result["tags"] == ["testing", "python"] + assert result["scope"] == "global" + assert result["timestamp"] == "1689201948123" + assert result["memory_id"] == "test_id_123" + + +def test_memory_entry_invalid_key(): + """Test with a key that does not start with 'memory:'.""" + lmem = { + "key": "notamemory:global:cat::123:hash", + "value": json.dumps({"data": "test", "tags": []}) + } + + result = _memory_entry(lmem) + assert result is None + + +def test_memory_entry_malformed_json_value(): + """Test with malformed/string value where JSON parsing fails.""" + valid_key = "memory:local:notes::1689201948123:a1b2c3d4e5f6" + # value is just a raw string, not JSON + lmem = { + "key": valid_key, + "value": "This is just a raw string without tags" + } + + result = _memory_entry(lmem) + + assert result is not None + assert result["data"] == "This is just a raw string without tags" + assert result["tags"] == [] # Falls back to empty tags list + assert result["category"] == "notes" + assert result["scope"] == "local" + + +def test_memory_entry_missing_fields(): + """Test gracefully handling dictionaries with missing keys.""" + # Missing 'value' and 'memory_id' + lmem1 = { + "key": "memory:global:ideas::123:hash" + } + result1 = _memory_entry(lmem1) + assert result1 is not None + assert result1["data"] == "" + assert result1["tags"] == [] + assert result1["memory_id"] == "" + + # Missing 'key' + lmem2 = { + "value": json.dumps({"data": "test", "tags": []}) + } + result2 = _memory_entry(lmem2) + assert result2 is None + + # Empty dict + result3 = _memory_entry({}) + assert result3 is None + + +@pytest.mark.parametrize( + "key, expected", + [ + ( + "memory:local:code::20240101T120000Z:abc123hash", + {"scope": "local", "category": "code", "timestamp": "20240101T120000Z"}, + ), + ( + "memory:global:general", + {"scope": "global", "category": "general", "timestamp": ""}, + ), + ( + "memory:local::20240101T120000Z:abc123hash", + {"scope": "local", "category": "", "timestamp": "20240101T120000Z"}, + ), + ( + "memory", + {"scope": "", "category": "", "timestamp": ""}, + ), + ( + "", + {"scope": "", "category": "", "timestamp": ""}, + ), + ( + None, + {"scope": "", "category": "", "timestamp": ""}, + ), + ( + "memory:global:category:with:colons::20240101T120000Z:abc123hash", + {"scope": "global", "category": "category", "timestamp": "20240101T120000Z"}, + ), + ( + "memory:global:general::20240101T120000Z", + {"scope": "global", "category": "general", "timestamp": "20240101T120000Z"}, + ), + ], + ids=[ + "happy_path", + "missing_timestamp_hash", + "missing_category", + "missing_scope_and_category", + "empty_string", + "invalid_type", + "extra_colons_in_category", + "missing_hash_but_has_timestamp", + ] +) +def test_parse_key(key, expected): + """Test _parse_key with various valid and invalid formats.""" + result = _parse_key(key) + assert result == expected + +def test_parse_memory_value_valid_json(): + raw_data = json.dumps({"data": "some data", "tags": ["tag1", "tag2"]}) + result = _parse_memory_value(raw_data) + assert result == {"data": "some data", "tags": ["tag1", "tag2"]} + + +def test_parse_memory_value_invalid_json(): + raw_data = "this is not json" + result = _parse_memory_value(raw_data) + assert result == {"data": "this is not json", "tags": []} + + +def test_parse_memory_value_type_error(): + raw_data = 12345 + result = _parse_memory_value(raw_data) # type: ignore[arg-type] + assert result == {"data": 12345, "tags": []} + + +def test_parse_memory_value_non_dict_json(): + raw_data = '"just a string"' + result = _parse_memory_value(raw_data) + assert result == "just a string" diff --git a/scripts/README.md b/scripts/README.md index 0139e8d3..4086a932 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -60,26 +60,26 @@ These tools are used to benchmark the prompt classifier and extract datasets fro --- -## 4. Integration Test Suite (Root Directory) +## 4. Integration Test Suite -The integration test suite is located in the root directory. Tests are categorized below based on their primary function: +The integration test suite is located in the `tests/` and `scripts/` directories. Tests are categorized below based on their primary function: ### Circuit Breaker Tests -- **`test_circuit_breaker.py`**: Unit/integration tests for the dual circuit breaker (`router/circuit_breaker.py`), covering independent Google/Vendor tiers and probe-granting logic. -- **`verify_breaker.py`**: Sanity verification check for the circuit breaker. -- **`test_a2_verify.py`**: Quick sanity integration check for the agy proxy circuit breaker. +- **`tests/test_circuit_breaker.py`**: Unit/integration tests for the dual circuit breaker (`router/circuit_breaker.py`), covering independent Google/Vendor tiers and probe-granting logic. +- **`scripts/verification/verify_breaker.py`**: Sanity verification check for the circuit breaker. +- **`tests/test_a2_verify.py`**: Quick sanity integration check for the agy proxy circuit breaker. ### Classifier Tests -- **`test_classifier_accuracy.py`**: Accuracy evaluation suite covering 25 system prompts. +- **`tests/test_classifier_accuracy.py`**: Accuracy evaluation suite covering 25 system prompts. ### Routing & Proxy Tests -- **`test_agy_tiers.py`**: Validates `agy` proxy model tier routing. -- **`test_antigravity.py`**: Tests the connection to the host Antigravity CLI daemon (`agentapi`). +- **`tests/test_agy_tiers.py`**: Validates `agy` proxy model tier routing. +- **`tests/test_antigravity.py`**: Tests the connection to the host Antigravity CLI daemon (`agentapi`). ### Performance & Monitoring Tests -- **`test_stream_latency.py`**: Measures Time-To-First-Token (TTFT) and token generation speed. +- **`tests/test_stream_latency.py`**: Measures Time-To-First-Token (TTFT) and token generation speed. ### Simulation Tests -- **`test_agy_behavior.py`**: Asserts the behavior of the `agy` CLI client under quota limits. -- **`test_quota_reset.sh`**: Simulates/triggers quota reset conditions. -- **`watch_quota.sh`**: Watch/polling script for observing quota status. +- **`tests/test_agy_behavior.py`**: Asserts the behavior of the `agy` CLI client under quota limits. +- **`scripts/test_quota_reset.sh`**: Simulates/triggers quota reset conditions. +- **`scripts/watch_quota.sh`**: Watch/polling script for observing quota status. diff --git a/scripts/benchmark_tokens.py b/scripts/benchmark_tokens.py new file mode 100644 index 00000000..d701b4eb --- /dev/null +++ b/scripts/benchmark_tokens.py @@ -0,0 +1,76 @@ +import sys +import os +from pathlib import Path + +# Set CONFIG_PATH and ROUTER_API_KEY for import +os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "router" / "config.yaml") +os.environ["ROUTER_API_KEY"] = "local-token" +# Add the parent directory and the router directory to the path +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "router")) + +from router.main import estimate_prompt_tokens + +def verify_accuracy(): + """Benchmarking utility to verify token estimation accuracy across content types.""" + # Test cases inspired by the problem description + test_cases = [ + { + "name": "English prose", + "content": "This is a standard English prose sentence intended to evaluate the accuracy of the token estimation heuristic for typical content. " * 5, + "actual_tokens": 110, + }, + { + "name": "Python code", + "content": """ +def calculate_factorial(n): + if n == 0: + return 1 + else: + return n * calculate_factorial(n-1) + +for i in range(10): + print(f"Factorial of {i} is {calculate_factorial(i)}") +""" * 3, + "actual_tokens": 150, + }, + { + "name": "CJK text", + "content": "这是一个测试，用于验证中文字符的令牌估算逻辑。它应该比字符计数更准确。" * 5, + "actual_tokens": 60, + }, + { + "name": "Whitespace-padded JSON", + "content": '{\n "key": "value",\n "nested": {\n "inner": "data"\n }\n}\n' * 5, + "actual_tokens": 60, + }, + { + "name": "Emoji", + "content": "🚀🔥-🤖✨-📈💎-🚨🛠️-🌐" * 5, + "actual_tokens": 25, + } + ] + + print(f"{'Case':<25} | {'Actual':<7} | {'Estimated':<9} | {'Error':<7}") + print("-" * 55) + + all_passed = True + for case in test_cases: + body = {"messages": [{"content": case["content"]}]} + est = estimate_prompt_tokens(body) - 50 # Subtract metadata overhead + error = abs(est - case["actual_tokens"]) / case["actual_tokens"] + print(f"{case['name']:<25} | {case['actual_tokens']:<7} | {est:<9} | {error:.1%}") + # Acceptance criteria: within ±25% for these rough heuristics + if error > 0.25: + print(f" --> FAILURE: {case['name']} error exceeds target threshold") + all_passed = False + + assert all_passed, "Token estimation accuracy benchmark failed" + +if __name__ == "__main__": + try: + verify_accuracy() + sys.exit(0) + except AssertionError as e: + print(f"\nERROR: {e}") + sys.exit(1) diff --git a/scripts/get_pr_status.py b/scripts/get_pr_status.py new file mode 100644 index 00000000..d65f65f1 --- /dev/null +++ b/scripts/get_pr_status.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +import subprocess +import json +import sys +from typing import Sequence + + +def run_cmd(argv: Sequence[str]) -> str: + """Runs a command and returns stripped stdout.""" + result = subprocess.run(argv, shell=False, capture_output=True, text=True, check=True, timeout=30) + return result.stdout.strip() + + +def get_pr_status(pr_id: str = "") -> None: + """Fetches and prints the status of a PR using gh CLI.""" + cmd = ["gh", "pr", "view"] + if pr_id: + cmd.append(pr_id) + cmd.extend(["--json", "state,reviewDecision,statusCheckRollup"]) + + try: + output = run_cmd(cmd) + data = json.loads(output) + + state = data.get("state") + review = data.get("reviewDecision") or "NONE" + checks = data.get("statusCheckRollup", []) + + # Summarize checks + success_count = 0 + total_count = len(checks) + for check in checks: + # gh CLI returns conclusion for CheckRun and state for StatusContext + conclusion = check.get("conclusion") or check.get("state") + if conclusion == "SUCCESS": + success_count += 1 + + print(f"PR Status: {state}") + print(f"Review Decision: {review}") + print(f"Checks: {success_count}/{total_count} passed") + + except subprocess.CalledProcessError as e: + print(f"Error: Failed to fetch PR status: {e.stderr.strip()}", file=sys.stderr) + sys.exit(1) + except json.JSONDecodeError: + print("Error: Failed to parse gh CLI output", file=sys.stderr) + sys.exit(1) + except Exception as e: + print(f"An unexpected error occurred: {e}", file=sys.stderr) + sys.exit(1) + + +def main(): + pr_id = sys.argv[1] if len(sys.argv) > 1 else "" + get_pr_status(pr_id) + + +if __name__ == "__main__": + main() diff --git a/host_agy_daemon.py b/scripts/host_agy_daemon.py similarity index 100% rename from host_agy_daemon.py rename to scripts/host_agy_daemon.py diff --git a/sync_gemini_token.py b/scripts/sync_gemini_token.py similarity index 100% rename from sync_gemini_token.py rename to scripts/sync_gemini_token.py diff --git a/test_quota_reset.sh b/scripts/test_quota_reset.sh similarity index 100% rename from test_quota_reset.sh rename to scripts/test_quota_reset.sh diff --git a/verify_breaker.py b/scripts/verification/verify_breaker.py similarity index 100% rename from verify_breaker.py rename to scripts/verification/verify_breaker.py diff --git a/watch_quota.sh b/scripts/watch_quota.sh similarity index 100% rename from watch_quota.sh rename to scripts/watch_quota.sh diff --git a/start-stack.sh b/start-stack.sh index c9b826dc..03bb1220 100755 --- a/start-stack.sh +++ b/start-stack.sh @@ -31,12 +31,6 @@ if [ -f "$ENV_FILE" ]; then fi # Ensure openssl is installed if we need to generate passwords/keys -if [ -z "$POSTGRES_PASSWORD" ] || [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ]; then - if ! command -v openssl &>/dev/null; then - echo "❌ Error: 'openssl' is required to generate secure random keys but was not found in PATH." - exit 1 - fi -fi if [ -z "$OPENROUTER_API_KEY" ]; then @@ -101,7 +95,7 @@ else echo "⚠️ Warning: Host agy daemon not responding on port 5005" fi -if [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ] || [ -z "$ROUTER_API_KEY" ]; then +if [ -z "$POSTGRES_PASSWORD" ] || [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ] || [ -z "$ROUTER_API_KEY" ] || [ -z "$MINIO_ROOT_USER" ] || [ -z "$MINIO_ROOT_PASSWORD" ]; then if ! command -v openssl &>/dev/null; then echo "❌ Error: 'openssl' is required to generate secure random keys but was not found in PATH." exit 1 @@ -141,12 +135,43 @@ if [ -z "$LITELLM_MASTER_KEY" ]; then exit 1 fi +if [ -z "$LANGFUSE_INIT_USER_PASSWORD" ]; then + LANGFUSE_INIT_USER_PASSWORD="$(openssl rand -hex 16)" + echo "LANGFUSE_INIT_USER_PASSWORD=\"$LANGFUSE_INIT_USER_PASSWORD\"" >> "$ENV_FILE" + echo "✓ Generated new LANGFUSE_INIT_USER_PASSWORD and saved to $ENV_FILE" +fi + +if [ -z "$REDIS_AUTH" ]; then + REDIS_AUTH="$(openssl rand -hex 16)" + echo "REDIS_AUTH=\"$REDIS_AUTH\"" >> "$ENV_FILE" + echo "✓ Generated new REDIS_AUTH and saved to $ENV_FILE" +fi + +if [ -z "$CLICKHOUSE_PASSWORD" ]; then + CLICKHOUSE_PASSWORD="$(openssl rand -hex 16)" + echo "CLICKHOUSE_PASSWORD=\"$CLICKHOUSE_PASSWORD\"" >> "$ENV_FILE" + echo "✓ Generated new CLICKHOUSE_PASSWORD and saved to $ENV_FILE" +fi + if [ -z "$ROUTER_API_KEY" ]; then ROUTER_API_KEY="$(openssl rand -hex 32)" echo "ROUTER_API_KEY=\"$ROUTER_API_KEY\"" >> "$ENV_FILE" echo "✓ Generated new ROUTER_API_KEY and saved to $ENV_FILE" fi +if [ -z "$MINIO_ROOT_USER" ]; then + MINIO_ROOT_USER="minio-$(openssl rand -hex 4)" + echo "MINIO_ROOT_USER=\"$MINIO_ROOT_USER\"" >> "$ENV_FILE" + echo "✓ Generated new MINIO_ROOT_USER and saved to $ENV_FILE" +fi + +if [ -z "$MINIO_ROOT_PASSWORD" ]; then + MINIO_ROOT_PASSWORD="$(openssl rand -hex 16)" + echo "MINIO_ROOT_PASSWORD=\"$MINIO_ROOT_PASSWORD\"" >> "$ENV_FILE" + echo "✓ Generated new MINIO_ROOT_PASSWORD and saved to $ENV_FILE" +fi + + # DYNAMIC_LITELLM_MASTER_KEY_PLACEHOLDER in router config is resolved at runtime from env @@ -255,7 +280,7 @@ setup_minio_buckets() { # Ensure mc alias points to the correct MinIO S3 API port (9002, not 9000) # The default 'local' alias in the MinIO image points to :9000 which is ClickHouse, # not MinIO. We must override it. - podman exec agent-router-pod-minio-s3 mc alias set local http://127.0.0.1:9002 minioadmin minioadmin 2>/dev/null + podman exec agent-router-pod-minio-s3 mc alias set local http://127.0.0.1:9002 "$MINIO_ROOT_USER" "$MINIO_ROOT_PASSWORD" 2>/dev/null # Create required buckets (idempotent) local BUCKETS=("langfuse-events" "proj-triage-gateway-id") @@ -356,7 +381,7 @@ if podman pod exists agent-router-pod 2>/dev/null; then fi render_pod_yaml() { - export WORKDIR HOME LITELLM_MASTER_KEY POSTGRES_PASSWORD NEXTAUTH_SECRET SALT ENCRYPTION_KEY + export WORKDIR HOME LITELLM_MASTER_KEY POSTGRES_PASSWORD NEXTAUTH_SECRET SALT ENCRYPTION_KEY LANGFUSE_INIT_USER_PASSWORD MINIO_ROOT_USER MINIO_ROOT_PASSWORD python3 - "$WORKDIR/pod.yaml" <<'PY' import os, sys, urllib.parse uid = os.getuid() @@ -371,11 +396,14 @@ placeholders = [ "NEXTAUTH_SECRET_PLACEHOLDER", "SALT_PLACEHOLDER", "ENCRYPTION_KEY_PLACEHOLDER", - "postgres-password-***" + "postgres-password-***", + "MINIO_USER_PLACEHOLDER", + "MINIO_PASSWORD_PLACEHOLDER" + "LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER" ] for ph in placeholders: if ph not in text: - sys.stderr.write(f"Error: Required placeholder '{ph}' not found in pod.yaml\n") + sys.stderr.write(f"Error: Required placeholder '{ph}' not found in pod.yaml. Ensure you are using the latest version of the template.\n") sys.exit(1) text = text.replace("/home/gpav/Vrac/LAB/AI/LLM-Routing", os.environ["WORKDIR"]) text = text.replace("/home/gpav/", os.environ["HOME"] + "/") @@ -388,6 +416,9 @@ text = text.replace("postgres-password-***", os.environ["POSTGRES_PASSWORD"]) text = text.replace("NEXTAUTH_SECRET_PLACEHOLDER", os.environ["NEXTAUTH_SECRET"]) text = text.replace("SALT_PLACEHOLDER", os.environ["SALT"]) text = text.replace("ENCRYPTION_KEY_PLACEHOLDER", os.environ["ENCRYPTION_KEY"]) +text = text.replace("MINIO_USER_PLACEHOLDER", os.environ["MINIO_ROOT_USER"]) +text = text.replace("MINIO_PASSWORD_PLACEHOLDER", os.environ["MINIO_ROOT_PASSWORD"]) +text = text.replace("LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER", os.environ["LANGFUSE_INIT_USER_PASSWORD"]) sys.stdout.write(text) PY } diff --git a/test_memory_mcp.py b/test_memory_mcp.py deleted file mode 100644 index 0194dca9..00000000 --- a/test_memory_mcp.py +++ /dev/null @@ -1,167 +0,0 @@ -#!/usr/bin/env python3 -""" -Tests for memory_mcp.py -""" -import sys -import json -import pytest -from router.memory_mcp import _memory_entry, _parse_key, _parse_memory_value - -def test_memory_entry_happy_path(): - """Test correctly formatted and complete memory entry.""" - valid_key = "memory:global:project_standards::1689201948123:a1b2c3d4e5f6" - valid_value = json.dumps({"data": "Use pytest for all tests", "tags": ["testing", "python"]}) - lmem = { - "key": valid_key, - "value": valid_value, - "memory_id": "test_id_123" - } - - result = _memory_entry(lmem) - - assert result is not None - assert result["key"] == valid_key - assert result["category"] == "project_standards" - assert result["data"] == "Use pytest for all tests" - assert result["tags"] == ["testing", "python"] - assert result["scope"] == "global" - assert result["timestamp"] == "1689201948123" - assert result["memory_id"] == "test_id_123" - -def test_memory_entry_invalid_key(): - """Test with a key that does not start with 'memory:'.""" - lmem = { - "key": "notamemory:global:cat::123:hash", - "value": json.dumps({"data": "test", "tags": []}) - } - - result = _memory_entry(lmem) - assert result is None - -def test_memory_entry_malformed_json_value(): - """Test with malformed/string value where JSON parsing fails.""" - valid_key = "memory:local:notes::1689201948123:a1b2c3d4e5f6" - # value is just a raw string, not JSON - lmem = { - "key": valid_key, - "value": "This is just a raw string without tags" - } - - result = _memory_entry(lmem) - - assert result is not None - assert result["data"] == "This is just a raw string without tags" - assert result["tags"] == [] # Falls back to empty tags list - assert result["category"] == "notes" - assert result["scope"] == "local" - -def test_memory_entry_missing_fields(): - """Test gracefully handling dictionaries with missing keys.""" - # Missing 'value' and 'memory_id' - lmem1 = { - "key": "memory:global:ideas::123:hash" - } - result1 = _memory_entry(lmem1) - assert result1 is not None - assert result1["data"] == "" - assert result1["tags"] == [] - assert result1["memory_id"] == "" - - # Missing 'key' - lmem2 = { - "value": json.dumps({"data": "test", "tags": []}) - } - result2 = _memory_entry(lmem2) - assert result2 is None - - # Empty dict - result3 = _memory_entry({}) - assert result3 is None - -def test_parse_key_happy_path(): - """Test full standard key structure""" - key = "memory:local:code::20240101T120000Z:abc123hash" - result = _parse_key(key) - assert result == { - "scope": "local", - "category": "code", - "timestamp": "20240101T120000Z" - } - -def test_parse_key_missing_timestamp_hash(): - """Test key without the :: delimiter section""" - key = "memory:global:general" - result = _parse_key(key) - assert result == { - "scope": "global", - "category": "general", - "timestamp": "" - } - -def test_parse_key_missing_category(): - """Test key with missing category""" - key = "memory:local::20240101T120000Z:abc123hash" - result = _parse_key(key) - # The split(":") on "memory:local" results in ["memory", "local"] length 2 - # So category should be "" - assert result == { - "scope": "local", - "category": "", - "timestamp": "20240101T120000Z" - } - -def test_parse_key_missing_scope_and_category(): - """Test minimal key prefix""" - key = "memory" - result = _parse_key(key) - assert result == { - "scope": "", - "category": "", - "timestamp": "" - } - -def test_parse_key_empty_string(): - """Test completely empty string""" - key = "" - result = _parse_key(key) - assert result == { - "scope": "", - "category": "", - "timestamp": "" - } - -def test_parse_key_invalid_type(): - """Test handling of an invalid type that triggers the exception branch""" - key = None - result = _parse_key(key) - assert result == { - "scope": "", - "category": "", - "timestamp": "" - } - -def test_parse_memory_value_valid_json(): - raw_data = json.dumps({"data": "some data", "tags": ["tag1", "tag2"]}) - result = _parse_memory_value(raw_data) - assert result == {"data": "some data", "tags": ["tag1", "tag2"]} - -def test_parse_memory_value_invalid_json(): - raw_data = "this is not json" - result = _parse_memory_value(raw_data) - assert result == {"data": "this is not json", "tags": []} - -def test_parse_memory_value_type_error(): - # json.loads will raise TypeError if given something that isn't str, bytes, or bytearray - raw_data = 12345 - result = _parse_memory_value(raw_data) # type: ignore[arg-type] - assert result == {"data": 12345, "tags": []} - -def test_parse_memory_value_non_dict_json(): - # If the input is valid JSON but not a dictionary, it currently returns the parsed non-dict value, - # which violates the dict return type annotation and can cause downstream KeyErrors/TypeErrors. - raw_data = '"just a string"' - result = _parse_memory_value(raw_data) - assert result == "just a string" - -if __name__ == "__main__": - sys.exit(pytest.main(["-v", __file__])) diff --git a/test_a2_verify.py b/tests/test_a2_verify.py similarity index 100% rename from test_a2_verify.py rename to tests/test_a2_verify.py diff --git a/test_agy_behavior.py b/tests/test_agy_behavior.py similarity index 96% rename from test_agy_behavior.py rename to tests/test_agy_behavior.py index 6dade8f6..37eb1dd1 100644 --- a/test_agy_behavior.py +++ b/tests/test_agy_behavior.py @@ -34,4 +34,5 @@ async def test(): if "RESOURCE_EXHAUSTED" in line or "quota" in line.lower(): print(f" {line.rstrip()}") -asyncio.run(test()) \ No newline at end of file +if __name__ == "__main__": + asyncio.run(test()) \ No newline at end of file diff --git a/test_agy_tiers.py b/tests/test_agy_tiers.py similarity index 100% rename from test_agy_tiers.py rename to tests/test_agy_tiers.py diff --git a/test_antigravity.py b/tests/test_antigravity.py similarity index 100% rename from test_antigravity.py rename to tests/test_antigravity.py diff --git a/test_atomic_write.py b/tests/test_atomic_write.py similarity index 100% rename from test_atomic_write.py rename to tests/test_atomic_write.py diff --git a/test_check_http_endpoint.py b/tests/test_check_http_endpoint.py similarity index 100% rename from test_check_http_endpoint.py rename to tests/test_check_http_endpoint.py diff --git a/test_circuit_breaker.py b/tests/test_circuit_breaker.py similarity index 100% rename from test_circuit_breaker.py rename to tests/test_circuit_breaker.py diff --git a/test_classifier_accuracy.py b/tests/test_classifier_accuracy.py similarity index 100% rename from test_classifier_accuracy.py rename to tests/test_classifier_accuracy.py diff --git a/test_compute_free_model_score.py b/tests/test_compute_free_model_score.py similarity index 100% rename from test_compute_free_model_score.py rename to tests/test_compute_free_model_score.py diff --git a/tests/test_get_pr_status.py b/tests/test_get_pr_status.py new file mode 100644 index 00000000..93885a14 --- /dev/null +++ b/tests/test_get_pr_status.py @@ -0,0 +1,84 @@ +import pytest +import subprocess +import json +from unittest.mock import patch, MagicMock +from scripts.get_pr_status import run_cmd, get_pr_status + +def test_run_cmd_success(): + output = run_cmd(["echo", "hello"]) + assert output == "hello" + +def test_run_cmd_strips_whitespace(): + output = run_cmd(["echo", " hello "]) + assert output == "hello" + +def test_run_cmd_error(): + with pytest.raises(subprocess.CalledProcessError): + run_cmd(["false"]) + +@patch("scripts.get_pr_status.subprocess.run") +def test_run_cmd_timeout(mock_run): + # run_cmd has a 30s timeout. We mock subprocess.run to raise it immediately. + mock_run.side_effect = subprocess.TimeoutExpired(["sleep", "0.1"], 30) + with pytest.raises(subprocess.TimeoutExpired): + run_cmd(["sleep", "0.1"]) + +@patch("scripts.get_pr_status.run_cmd") +def test_get_pr_status_success(mock_run_cmd, capsys): + mock_data = { + "state": "OPEN", + "reviewDecision": "APPROVED", + "statusCheckRollup": [ + {"conclusion": "SUCCESS", "name": "test1"}, + {"state": "SUCCESS", "name": "test2"}, + {"conclusion": "FAILURE", "name": "test3"} + ] + } + mock_run_cmd.return_value = json.dumps(mock_data) + + get_pr_status("123") + + captured = capsys.readouterr() + assert "PR Status: OPEN" in captured.out + assert "Review Decision: APPROVED" in captured.out + assert "Checks: 2/3 passed" in captured.out + mock_run_cmd.assert_called_once_with(["gh", "pr", "view", "123", "--json", "state,reviewDecision,statusCheckRollup"]) + +@patch("scripts.get_pr_status.run_cmd") +def test_get_pr_status_no_id(mock_run_cmd, capsys): + mock_data = { + "state": "MERGED", + "reviewDecision": None, + "statusCheckRollup": [] + } + mock_run_cmd.return_value = json.dumps(mock_data) + + get_pr_status() + + captured = capsys.readouterr() + assert "PR Status: MERGED" in captured.out + assert "Review Decision: NONE" in captured.out + assert "Checks: 0/0 passed" in captured.out + mock_run_cmd.assert_called_once_with(["gh", "pr", "view", "--json", "state,reviewDecision,statusCheckRollup"]) + +@patch("scripts.get_pr_status.run_cmd") +def test_get_pr_status_error(mock_run_cmd, capsys): + mock_run_cmd.side_effect = subprocess.CalledProcessError(1, ["gh"], stderr="gh not found") + + with pytest.raises(SystemExit) as e: + get_pr_status("123") + + assert e.value.code == 1 + captured = capsys.readouterr() + assert "Error: Failed to fetch PR status: gh not found" in captured.err + +@patch("scripts.get_pr_status.run_cmd") +def test_get_pr_status_invalid_json(mock_run_cmd, capsys): + mock_run_cmd.return_value = "invalid json" + + with pytest.raises(SystemExit) as e: + get_pr_status("123") + + assert e.value.code == 1 + captured = capsys.readouterr() + assert "Error: Failed to parse gh CLI output" in captured.err diff --git a/tests/test_host_agy_daemon.py b/tests/test_host_agy_daemon.py new file mode 100644 index 00000000..61f84e86 --- /dev/null +++ b/tests/test_host_agy_daemon.py @@ -0,0 +1,490 @@ +import asyncio +import json +import os +import socket +import threading +import urllib.error +import urllib.request +from unittest.mock import AsyncMock + +import pytest + +import sys +from pathlib import Path + +# Dynamic project root discovery +root = Path(__file__).resolve() +while root.parent != root and not (root / ".git").exists(): + root = root.parent +sys.path.insert(0, str(root)) +sys.path.insert(0, str(root / "scripts")) + +import host_agy_daemon + +def find_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('127.0.0.1', 0)) + return s.getsockname()[1] + +@pytest.fixture +def daemon_server(): + port = find_free_port() + host_agy_daemon.PORT = port + + server = host_agy_daemon.ThreadingHTTPServer(('127.0.0.1', port), host_agy_daemon.AgyDaemonHandler) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + + yield f"http://127.0.0.1:{port}" + + server.shutdown() + server.server_close() + server_thread.join() + +def test_get_last_conversation_id(monkeypatch, tmp_path): + cache_file = tmp_path / "last_conversations.json" + cache_file.write_text(json.dumps({"/fake/cwd": "conv_123"})) + + monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", str(cache_file)) + monkeypatch.setattr(host_agy_daemon.os, "getcwd", lambda: "/fake/cwd") + + assert host_agy_daemon.get_last_conversation_id() == "conv_123" + + monkeypatch.setattr(host_agy_daemon.os, "getcwd", lambda: "/other/cwd") + assert host_agy_daemon.get_last_conversation_id() is None + +def test_get_last_conversation_id_no_file(monkeypatch): + monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", "/does/not/exist.json") + assert host_agy_daemon.get_last_conversation_id() is None + +def test_get_last_conversation_id_invalid_json(monkeypatch, tmp_path): + cache_file = tmp_path / "last_conversations.json" + cache_file.write_text("invalid json") + + monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", str(cache_file)) + assert host_agy_daemon.get_last_conversation_id() is None + +def test_get_last_conversation_id_io_error(monkeypatch): + monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", "/fake/cache.json") + monkeypatch.setattr(host_agy_daemon.os.path, "exists", lambda x: True) + def mock_open_err(*args, **kwargs): + raise IOError("permission denied") + monkeypatch.setattr("builtins.open", mock_open_err) + assert host_agy_daemon.get_last_conversation_id() is None + +def test_daemon_post_404(daemon_server): + req = urllib.request.Request(f"{daemon_server}/invalid", method="POST") + with pytest.raises(urllib.error.HTTPError) as exc: + urllib.request.urlopen(req) + assert exc.value.code == 404 + +def test_daemon_post_stream_false(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": False, "conversation_id": "conv_abc", "model_override": "gpt-4"}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + assert args == (host_agy_daemon.AGY_BINARY, "--conversation", "conv_abc", "--print", "test prompt") + assert kwargs.get("env", {}).get("CASCADE_DEFAULT_MODEL_OVERRIDE") == "gpt-4" + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.wait = AsyncMock() + + if "stdout" in kwargs: + with open(kwargs["stdout"].name, "w") as f: + f.write("mocked stdout output") + if "stderr" in kwargs: + with open(kwargs["stderr"].name, "w") as f: + f.write("mocked stderr output") + + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: "last_conv_456") + + with urllib.request.urlopen(req) as resp: + data = json.loads(resp.read().decode()) + + assert data["returncode"] == 0 + assert data["stdout"] == "mocked stdout output" + assert data["stderr"] == "mocked stderr output" + assert data["conversation_id"] == "last_conv_456" + +def test_daemon_post_stream_false_timeout(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": False, "timeout": 0.1}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + # Make wait take longer than timeout + async def slow_wait(): + await asyncio.sleep(0.5) + mock_proc.wait = slow_wait + # Make kill synchronous + mock_proc.kill = lambda: None + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + + with urllib.request.urlopen(req) as resp: + data = json.loads(resp.read().decode()) + + assert data["returncode"] == -1 + assert data["stderr"] == "TIMEOUT" + +def test_daemon_post_stream_true(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": True, "model_override": "test-model"}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + assert args == (host_agy_daemon.AGY_BINARY, "--print", "test prompt") + assert kwargs.get("env", {}).get("CASCADE_DEFAULT_MODEL_OVERRIDE") == "test-model" + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.wait = AsyncMock() + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: "last_conv_456") + + read_calls = 0 + def mock_read(fd, n): + nonlocal read_calls + if read_calls == 0: + read_calls += 1 + return b"token1\r\n" + elif read_calls == 1: + read_calls += 1 + return b"token2\r\n" + return b"" + + monkeypatch.setattr(host_agy_daemon.os, "read", mock_read) + + with urllib.request.urlopen(req) as resp: + content = resp.read().decode().strip() + lines = content.split("\n") + + assert len(lines) == 3 + assert json.loads(lines[0]) == {"type": "token", "content": "token1\n"} + assert json.loads(lines[1]) == {"type": "token", "content": "token2\n"} + assert json.loads(lines[2]) == {"type": "status", "returncode": 0, "conversation_id": "last_conv_456"} + +def test_daemon_post_stream_true_exec_error(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": True}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + raise Exception("exec failed") + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + + with urllib.request.urlopen(req) as resp: + content = resp.read().decode().strip() + lines = content.split("\n") + + assert len(lines) == 1 + assert json.loads(lines[0]) == {"type": "status", "returncode": -1, "stderr": "exec failed"} + +def test_daemon_post_stream_true_timeout(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": True, "timeout": 0.1}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + async def slow_wait(): + await asyncio.sleep(0.5) + mock_proc.wait = slow_wait + # Make kill synchronous + mock_proc.kill = lambda: None + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: None) + + read_calls = 0 + def mock_read(fd, n): + nonlocal read_calls + if read_calls == 0: + read_calls += 1 + return b"token1\n" + return b"" + + monkeypatch.setattr(host_agy_daemon.os, "read", mock_read) + + with urllib.request.urlopen(req) as resp: + content = resp.read().decode().strip() + lines = content.split("\n") + + assert len(lines) == 2 + assert json.loads(lines[0]) == {"type": "token", "content": "token1\n"} + assert json.loads(lines[1]) == {"type": "status", "returncode": -1, "conversation_id": None} + +def test_log_message_silenced(): + # Instantiate the class, bypassing BaseHTTPRequestHandler.__init__ + handler = host_agy_daemon.AgyDaemonHandler.__new__(host_agy_daemon.AgyDaemonHandler) + # Shouldn't raise any error + handler.log_message("format %s", "arg") + +def test_run_server_interrupt(monkeypatch): + # Mock serve_forever to raise KeyboardInterrupt + def mock_serve_forever(self): + raise KeyboardInterrupt() + + monkeypatch.setattr(host_agy_daemon.ThreadingHTTPServer, "serve_forever", mock_serve_forever) + + # Track if server_close was called + close_called = False + def mock_server_close(self): + nonlocal close_called + close_called = True + + monkeypatch.setattr(host_agy_daemon.ThreadingHTTPServer, "server_close", mock_server_close) + + # Should not raise exception + host_agy_daemon.run_server() + assert close_called + +def test_daemon_post_stream_false_no_model_override(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + assert "CASCADE_DEFAULT_MODEL_OVERRIDE" not in kwargs.get("env", {}) + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.wait = AsyncMock() + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + monkeypatch.setattr(host_agy_daemon.os.environ, "copy", lambda: {"CASCADE_DEFAULT_MODEL_OVERRIDE": "old-model"}) + + with urllib.request.urlopen(req) as resp: + data = json.loads(resp.read().decode()) + + assert data["returncode"] == 0 + +def test_daemon_post_stream_true_read_oserror(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": True}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.wait = AsyncMock() + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + + def mock_read(fd, n): + raise OSError("read error") + + monkeypatch.setattr(host_agy_daemon.os, "read", mock_read) + + with urllib.request.urlopen(req) as resp: + content = resp.read().decode().strip() + lines = content.split("\n") + + assert len(lines) == 1 + assert json.loads(lines[0])["type"] == "status" + +def test_daemon_post_stream_true_timeout_kill_fail(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": True, "timeout": 0.1}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + async def slow_wait(): + await asyncio.sleep(0.5) + mock_proc.wait = slow_wait + def mock_kill(): + raise Exception("kill failed") + mock_proc.kill = mock_kill + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + monkeypatch.setattr(host_agy_daemon.os, "read", lambda fd, n: b"") + monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: None) + + with urllib.request.urlopen(req) as resp: + content = resp.read().decode().strip() + lines = content.split("\n") + + assert len(lines) == 1 + assert json.loads(lines[0]) == {"type": "status", "returncode": -1, "conversation_id": None} + +def test_daemon_post_stream_true_wait_exception(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": True}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + async def mock_wait(): + raise Exception("wait failed") + mock_proc.wait = mock_wait + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + monkeypatch.setattr(host_agy_daemon.os, "read", lambda fd, n: b"") + monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: None) + + with urllib.request.urlopen(req) as resp: + content = resp.read().decode().strip() + lines = content.split("\n") + + assert len(lines) == 1 + assert json.loads(lines[0]) == {"type": "status", "returncode": -1, "conversation_id": None} + +def test_daemon_post_stream_false_timeout_kill_fail(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": False, "timeout": 0.1}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + async def slow_wait(): + await asyncio.sleep(0.5) + mock_proc.wait = slow_wait + def mock_kill(): + raise Exception("kill failed") + mock_proc.kill = mock_kill + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + + with urllib.request.urlopen(req) as resp: + data = json.loads(resp.read().decode()) + + assert data["returncode"] == -1 + assert data["stderr"] == "TIMEOUT" + +def test_daemon_post_stream_false_wait_exception(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + async def mock_wait(): + raise Exception("wait failed") + mock_proc.wait = mock_wait + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + + with urllib.request.urlopen(req) as resp: + data = json.loads(resp.read().decode()) + + assert data["returncode"] == -1 + +def test_daemon_post_stream_false_file_read_error(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.wait = AsyncMock() + + # Corrupt the temp files to cause read exceptions + os.unlink(kwargs["stdout"].name) + os.unlink(kwargs["stderr"].name) + + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + + with urllib.request.urlopen(req) as resp: + data = json.loads(resp.read().decode()) + + assert data["returncode"] == 0 + assert data["stdout"] == "" + assert data["stderr"] == "" + +def test_daemon_post_stream_false_unlink_error(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.wait = AsyncMock() + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + + def mock_unlink(path): + raise Exception("unlink failed") + + monkeypatch.setattr(host_agy_daemon.os, "unlink", mock_unlink) + + with urllib.request.urlopen(req) as resp: + data = json.loads(resp.read().decode()) + + assert data["returncode"] == 0 + +def test_daemon_post_stream_true_with_conversation(daemon_server, monkeypatch): + req = urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps({"prompt": "test prompt", "stream": True, "conversation_id": "conv_789"}).encode("utf-8"), + headers={"Content-Type": "application/json"} + ) + + async def mock_exec(*args, **kwargs): + assert args == (host_agy_daemon.AGY_BINARY, "--conversation", "conv_789", "--print", "test prompt") + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.wait = AsyncMock() + return mock_proc + + monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) + monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: "conv_789") + monkeypatch.setattr(host_agy_daemon.os, "read", lambda fd, n: b"") + + with urllib.request.urlopen(req) as resp: + content = resp.read().decode().strip() + lines = content.split("\n") + + assert len(lines) == 1 + assert json.loads(lines[0]) == {"type": "status", "returncode": 0, "conversation_id": "conv_789"} diff --git a/test_map_tool_to_category.py b/tests/test_map_tool_to_category.py similarity index 100% rename from test_map_tool_to_category.py rename to tests/test_map_tool_to_category.py diff --git a/test_models_proxy.py b/tests/test_models_proxy.py similarity index 100% rename from test_models_proxy.py rename to tests/test_models_proxy.py diff --git a/test_pie_chart_gradient.py b/tests/test_pie_chart_gradient.py similarity index 100% rename from test_pie_chart_gradient.py rename to tests/test_pie_chart_gradient.py diff --git a/tests/test_read_annotations_async.py b/tests/test_read_annotations_async.py new file mode 100644 index 00000000..b401bc8b --- /dev/null +++ b/tests/test_read_annotations_async.py @@ -0,0 +1,128 @@ +import pytest +from unittest.mock import patch, AsyncMock, MagicMock +import copy +import sys +import os +import json +import asyncio + +# Ensure the root directory is in the path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +import router.main +from router.main import _read_annotations_async + +@pytest.fixture(autouse=True) +def clear_cache(): + router.main._annotations_cache.clear() + yield + +@pytest.mark.asyncio +async def test_read_annotations_async_initial_read(): + fake_path = "/tmp/annotations.json" + fake_data = {"annotation1": "data1"} + + # Mock aiofiles.open + mock_file = AsyncMock() + mock_file.read.return_value = '{"annotation1": "data1"}' + + mock_context_manager = MagicMock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_file) + mock_context_manager.__aexit__ = AsyncMock(return_value=False) + + mock_aiofiles_open = MagicMock(return_value=mock_context_manager) + + with patch("os.path.getmtime", return_value=100.0) as mock_getmtime, \ + patch("aiofiles.open", mock_aiofiles_open) as mock_open: + + result = await _read_annotations_async(fake_path) + + mock_getmtime.assert_called_once_with(fake_path) + mock_open.assert_called_once_with(fake_path, "r", encoding="utf-8") + assert result == fake_data + + # Verify cache is populated + assert fake_path in router.main._annotations_cache + assert router.main._annotations_cache[fake_path]["mtime"] == 100.0 + assert router.main._annotations_cache[fake_path]["data"] == fake_data + +@pytest.mark.asyncio +async def test_read_annotations_async_cache_hit(): + fake_path = "/tmp/annotations.json" + fake_data = {"annotation1": "data1"} + + # Pre-populate cache + router.main._annotations_cache[fake_path] = {"mtime": 100.0, "data": fake_data} + + # Mock aiofiles.open (should NOT be called) + mock_aiofiles_open = MagicMock() + + with patch("os.path.getmtime", return_value=100.0) as mock_getmtime, \ + patch("aiofiles.open", mock_aiofiles_open) as mock_open: + + result = await _read_annotations_async(fake_path) + + mock_getmtime.assert_called_once_with(fake_path) + mock_open.assert_not_called() + assert result == fake_data + +@pytest.mark.asyncio +async def test_read_annotations_async_cache_invalidation(): + fake_path = "/tmp/annotations.json" + fake_data_old = {"annotation1": "data1"} + fake_data_new = {"annotation2": "data2"} + + # Pre-populate cache with old mtime + router.main._annotations_cache[fake_path] = {"mtime": 100.0, "data": fake_data_old} + + mock_file = AsyncMock() + mock_file.read.return_value = '{"annotation2": "data2"}' + + mock_context_manager = MagicMock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_file) + mock_context_manager.__aexit__ = AsyncMock(return_value=False) + + mock_aiofiles_open = MagicMock(return_value=mock_context_manager) + + with patch("os.path.getmtime", return_value=200.0) as mock_getmtime, \ + patch("aiofiles.open", mock_aiofiles_open) as mock_open: + + result = await _read_annotations_async(fake_path) + + mock_getmtime.assert_called_once_with(fake_path) + mock_open.assert_called_once_with(fake_path, "r", encoding="utf-8") + assert result == fake_data_new + + # Verify cache is updated + assert router.main._annotations_cache[fake_path]["mtime"] == 200.0 + assert router.main._annotations_cache[fake_path]["data"] == fake_data_new + +@pytest.mark.asyncio +async def test_read_annotations_async_deepcopy(): + fake_path = "/tmp/annotations.json" + fake_data = {"annotation1": {"nested": "value"}} + + # Pre-populate cache + router.main._annotations_cache[fake_path] = {"mtime": 100.0, "data": fake_data} + + with patch("os.path.getmtime", return_value=100.0): + # First read + result1 = await _read_annotations_async(fake_path) + + # Mutate the result + result1["annotation1"]["nested"] = "mutated" + + # Second read + result2 = await _read_annotations_async(fake_path) + + # Verify second read returns original data, not mutated + assert result2["annotation1"]["nested"] == "value" + assert router.main._annotations_cache[fake_path]["data"]["annotation1"]["nested"] == "value" + +@pytest.mark.asyncio +async def test_read_annotations_async_file_not_found(): + fake_path = "/tmp/annotations.json" + + with patch("os.path.getmtime", side_effect=FileNotFoundError): + with pytest.raises(FileNotFoundError): + await _read_annotations_async(fake_path) diff --git a/test_record_tool_usage.py b/tests/test_record_tool_usage.py similarity index 100% rename from test_record_tool_usage.py rename to tests/test_record_tool_usage.py diff --git a/test_src_badge.py b/tests/test_src_badge.py similarity index 100% rename from test_src_badge.py rename to tests/test_src_badge.py diff --git a/test_stream_latency.py b/tests/test_stream_latency.py similarity index 100% rename from test_stream_latency.py rename to tests/test_stream_latency.py diff --git a/test_sync_gemini_token.py b/tests/test_sync_gemini_token.py similarity index 100% rename from test_sync_gemini_token.py rename to tests/test_sync_gemini_token.py