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/branch_audit_plan.md b/.agents/branch_audit_plan.md new file mode 100644 index 00000000..f1a9374c --- /dev/null +++ b/.agents/branch_audit_plan.md @@ -0,0 +1,56 @@ +# Git Merges & Regressions Audit Report + +## Problem Overview +During the merging of concurrent pull requests, automated conflict-resolution bots (`jules[bot]`) silently discarded or reverted several code blocks and logic optimizations. + +To ensure complete stability, we performed a systematic, automated three-dot diffing audit (`git diff CURRENT_BRANCH...TARGET_BRANCH`) comparing our current workspace against every active and merged remote branch in the repository. + +--- + +## Branch-by-Branch Audit Checklist + +| Branch Name / PR | Status | Notes / Discrepancies Checked | +| :--- | :---: | :--- | +| `origin/cleanup/remove-unused-get-live-gemini-oauth-token` (PR #195) | **CLEAN** | All changes present in current codebase. | +| `origin/perf-optimize-gemini-oauth-token-7488735575783948165` (PR #158) | **CLEAN** | All changes present in current codebase. | +| `origin/refactor-record-tool-usage-17779903941173107649` (PR #157) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/get-pr-status-tests-18330387033569416831` (PR #193) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/perf-async-annotations-4499495305587811104` (PR #172) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/add-gemini-oauth-tests-17524029845400706758` (PR #163) | **CLEAN** | All changes present in current codebase. | +| `origin/dependabot/docker/berriai/litellm-v1.90.2` (PR #192) | **CLEAN** | Diff showed `v1.90.2` LiteLLM image version bump, which is already applied. | +| `origin/security/fix-memory-mcp-category-injection-14795006529015952965` (PR #173) | **CLEAN** | All changes present in current codebase. | +| `origin/test/add-read-annotations-sync-tests-12043640228306609857` (PR #168) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/test-parse-key-parameterized-14651890987031951450` (PR #169) | **CLEAN** | All changes present in current codebase. | +| `origin/perf/async-quota-exhausted-4722057800691922723` (PR #156) | **CLEAN** | All changes present in current codebase. | +| `origin/perf-benchmark-classifier-14678555128983342452` (PR #155) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/add-test-redis-from-url-12746122485451637611` (PR #174) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/test-is-memory-key-14668590644102848561` (PR #164) | **CLEAN** | All changes present in current codebase. | +| `origin/⚡/optimize-agy-log-read-1519383298987985129` (PR #170) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/add-oauth-status-tests-15591368961301252072` (PR #166) | **CLEAN** | All changes present in current codebase. | +| `origin/chore/test-get-goose-sessions-16558526271462662412` (PR #165) | **CLEAN** | All changes present in current codebase. | +| `origin/fix-http-client-proliferation-11317864999360875839` | **CLEAN** | HTTP client unification is fully present in our workspace. | +| `origin/fix/test-record-tool-usage-3580129709200353766` | **CLEAN** | Root-level test files were reorganized to `/tests`, fully preserved. | +| `origin/fix/secrets-relocation` (PR #184) | **CLEAN** | Checked out placeholders and updated `start-stack.sh`/`pod.yaml` securely. | +| `origin/perf/offload-aa-scores-sync-16551127323385849872` (PR #98) | **RESTORED** | **Regression Found**: The asynchronous scores loading optimization was lost on master. Re-implemented and verified. | + +--- + +## Detailed Regression Analysis & Fix: Offload AA Scores Loading + +### The Regression +* **File:** `router/main.py` +* **Issue:** `_load_aa_scores()` (which does a blocking synchronous JSON disk read of `aa_scores.json`) was called directly inside `compute_free_model_score()`. When resolving or syncing OpenRouter free models, this synchronous file read was executed sequentially in a loop, blocking the main event loop. +* **Merged Fix Lost:** The merged PR #98 had offloaded this operation to a background thread pool using `await asyncio.to_thread(_load_aa_scores)` outside the loops in `sync_adaptive_router_roster` and `get_best_free_model`, but this change was completely discarded during automated conflict resolution. + +### The Resolution +1. **`router/main.py`**: + * Restored `await asyncio.to_thread(_load_aa_scores)` inside `sync_adaptive_router_roster()` and `get_best_free_model()` if the cache is not yet loaded. + * Removed the blocking `_load_aa_scores()` call from inside `compute_free_model_score()`. +2. **`tests/test_compute_free_model_score.py`**: + * Updated the test cases to explicitly call `router_main._load_aa_scores()` within the test setups to align with the asynchronous offloading. + +--- + +## Verification Status +* **Test Suite:** Ran `PYTHONPATH=router:. pytest` +* **Result:** **176 tests passed** successfully (100% success rate). diff --git a/.agents/jules_regression_analysis.md b/.agents/jules_regression_analysis.md new file mode 100644 index 00000000..37b17229 --- /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 --porcelain +``` + +--- + +## 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..8b22f7f3 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, ClickHouse, and Valkey Cache image tags are explicitly pinned in `pod.yaml` — never use `:latest`. Check available tags with `skopeo list-tags` or registry hubs before upgrading. 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/debug-homepage.png b/debug-homepage.png new file mode 100644 index 00000000..8a45fdc2 Binary files /dev/null and b/debug-homepage.png differ 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/litellm/entrypoint.py b/litellm/entrypoint.py index 1b987365..df7a2ec8 100644 --- a/litellm/entrypoint.py +++ b/litellm/entrypoint.py @@ -98,6 +98,28 @@ def strptime(cls, date_str: str, fmt: str) -> original_datetime: datetime.datetime = RobustDatetime sys.stdout.flush() +# Register both RobustDatetime AND the original datetime with Prisma's +# singledispatch serializer. When entrypoint.py replaces datetime.datetime +# with RobustDatetime before Prisma loads, Prisma's own +# @serializer.register(datetime.datetime) ends up registering RobustDatetime. +# But database drivers (psycopg2) return the *original* C-level datetime +# instances, which no longer match. We must register both classes. +try: + from prisma.builder import serializer +except ImportError: + serializer = None + +if serializer is not None: + def _serialize_dt(dt): + """Serialize datetime to ISO8601 with timezone (UTC if naive).""" + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt.isoformat() + serializer.register(original_datetime, _serialize_dt) + serializer.register(RobustDatetime, _serialize_dt) + print("🩹 Registered original_datetime + RobustDatetime with Prisma serializer") +sys.stdout.flush() + # Start LiteLLM Proxy import litellm from litellm.proxy.proxy_cli import run_server diff --git a/pod.yaml b/pod.yaml index 9d3af258..b8ed35b8 100644 --- a/pod.yaml +++ b/pod.yaml @@ -4,7 +4,7 @@ metadata: name: agent-router-pod spec: containers: - # Valkey Cache 9.1.0-alpine — LiteLLM response cache (exact-match & semantic) + # Valkey Cache — LiteLLM response cache (exact-match & semantic) - command: - valkey-server - --protected-mode @@ -34,7 +34,7 @@ spec: - /app/entrypoint.py env: - name: DATABASE_URL - value: postgresql://postgres:***@127.0.0.1:5432/postgres + value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/postgres - name: STORE_MODEL_IN_DB value: 'True' - name: LITELLM_LOG @@ -42,10 +42,10 @@ spec: - name: LANGFUSE_HOST value: http://127.0.0.1:3001 - name: LITELLM_MASTER_KEY - value: sk-lit...33bf + value: LITELLM_MASTER_KEY_PLACEHOLDER - name: OLLAMA_API_KEY - value: 3cd542f8fca14fd08bedfd4f2ab36f9f.6G7Iukbvu1Keyi9x8eKckNEO - image: ghcr.io/berriai/litellm:v1.89.4 + value: OLLAMA_API_KEY_PLACEHOLDER + image: ghcr.io/berriai/litellm:v1.90.2 livenessProbe: exec: command: @@ -89,19 +89,19 @@ spec: - name: LITELLM_CONFIG_PATH value: /config/litellm_dir/config.yaml - name: DATABASE_URL - value: postgresql://postgres:***@127.0.0.1:5432/postgres + value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/postgres - name: DBUS_SESSION_BUS_ADDRESS value: unix:path=/run/user/1000/bus - name: LITELLM_MASTER_KEY - value: sk-lit...33bf + value: LITELLM_MASTER_KEY_PLACEHOLDER - name: LANGFUSE_PUBLIC_KEY - value: pk-lf-200bbf48-1d23-447d-a79a-8614106497e6 + value: LANGFUSE_PUBLIC_KEY_PLACEHOLDER - name: LANGFUSE_SECRET_KEY - value: sk-lf-...3b49 + value: LANGFUSE_SECRET_KEY_PLACEHOLDER - name: LANGFUSE_HOST value: http://127.0.0.1:3001 - name: OLLAMA_API_KEY - value: 3cd542f8fca14fd08bedfd4f2ab36f9f.6G7Iukbvu1Keyi9x8eKckNEO + value: OLLAMA_API_KEY_PLACEHOLDER - name: LOG_LEVEL value: info image: localhost/llm-triage-router:latest @@ -154,10 +154,10 @@ spec: - name: POSTGRES_USER value: postgres - name: POSTGRES_PASSWORD - value: postgres-password-*** + value: POSTGRES_PASSWORD_RAW_PLACEHOLDER - name: POSTGRES_DB value: langfuse - image: docker.io/pgvector/pgvector:pg18 + image: docker.io/pgvector/pgvector:0.8.3-pg18 livenessProbe: exec: command: @@ -188,7 +188,7 @@ spec: value: clickhouse - name: CLICKHOUSE_PASSWORD value: clickhouse - image: docker.io/clickhouse/clickhouse-server:26.5.1 + image: docker.io/clickhouse/clickhouse-server:26.6.1.1193 livenessProbe: exec: command: @@ -253,7 +253,7 @@ spec: # Langfuse Web — observability dashboard - env: - name: DATABASE_URL - value: postgresql://postgres:***@127.0.0.1:5432/langfuse + value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/langfuse - name: NEXTAUTH_SECRET value: NEXTAUTH_SECRET_PLACEHOLDER - name: NEXTAUTH_URL @@ -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,10 +307,10 @@ 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 + image: docker.io/langfuse/langfuse:3.202.1 livenessProbe: exec: command: @@ -337,7 +337,7 @@ spec: # Langfuse Worker — background job processor - env: - name: DATABASE_URL - value: postgresql://postgres:***@127.0.0.1:5432/langfuse + value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/langfuse - name: CLICKHOUSE_URL value: http://127.0.0.1:8123 - name: CLICKHOUSE_USER @@ -363,16 +363,16 @@ 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 + image: docker.io/langfuse/langfuse-worker:3.202.1 livenessProbe: exec: command: @@ -393,10 +393,10 @@ spec: - ":9001" env: - name: MINIO_ROOT_USER - value: minioadmin + value: MINIO_USER_PLACEHOLDER - name: MINIO_ROOT_PASSWORD - value: minioadmin - image: docker.io/minio/minio:latest + value: MINIO_PASSWORD_PLACEHOLDER + image: docker.io/minio/minio:RELEASE.2025-09-07T16-13-09Z livenessProbe: httpGet: path: /minio/health/live 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/agy_proxy.py b/router/agy_proxy.py index 5da0c335..4e911c42 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -20,6 +20,7 @@ """ import json +import aiofiles import logging import os import time @@ -127,7 +128,7 @@ async def _run_agy_print(client: httpx.AsyncClient, prompt: str, model_override: # Track the last log check time to avoid hammering the file _last_log_check: float = 0 -def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool: +async def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool: """ Detect quota exhaustion from agy subprocess results. @@ -151,11 +152,18 @@ def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool: _last_log_check = now log_path = os.path.expanduser("~/.gemini/antigravity-cli/cli.log") try: - if os.path.exists(log_path): - with open(log_path, "r") as f: - for line in f.readlines()[-5:]: - if "RESOURCE_EXHAUSTED" in line or "code 429" in line: - return True + async with aiofiles.open(log_path, "rb") as f: + try: + await f.seek(0, 2) + file_size = await f.tell() + await f.seek(max(0, file_size - 1024)) + except OSError: + pass + content_bytes = await f.read() + content = content_bytes.decode("utf-8", errors="ignore") + for line in content.splitlines()[-5:]: + if "RESOURCE_EXHAUSTED" in line or "code 429" in line: + return True except Exception: pass # Empty stdout+stderr with rc=0 strongly suggests quota exhaustion @@ -346,8 +354,9 @@ async def try_agy_proxy(prompt: str, messages: list = None, rc = 0 if raw_rc is None else raw_rc raw_stderr = first_data.get("stderr", "") stderr_content = "" if raw_stderr is None else raw_stderr - if _is_quota_exhausted(rc, "", stderr_content) or rc != 0: - if _is_quota_exhausted(rc, "", stderr_content): + is_exhausted = await _is_quota_exhausted(rc, "", stderr_content) + if is_exhausted or rc != 0: + if is_exhausted: tier_breaker.record_failure() if cooldown_persistence is not None: try: @@ -420,7 +429,7 @@ async def token_generator(stream_resp, httpx_client, initial_line, current_conv_ last_conv_id = result_conv_id # Check for quota exhaustion - if _is_quota_exhausted(returncode, stdout, stderr): + if await _is_quota_exhausted(returncode, stdout, stderr): tier_breaker.record_failure() if cooldown_persistence is not None: try: diff --git a/router/main.py b/router/main.py index 6c1c9a2d..d9292f99 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, round(total) + 50) async def sync_cooldowns_from_valkey() -> None: @@ -487,6 +530,8 @@ async def sync_adaptive_router_roster(master_key: str): except Exception as e: logger.warning(f"Failed to fetch OpenRouter models: {e}") return + if not _AA_SCORES_LOADED: + await asyncio.to_thread(_load_aa_scores) free_models = [] model_contexts = {} model_supported_params = {} @@ -1407,7 +1452,6 @@ def _load_aa_scores(): def compute_free_model_score(m: dict) -> float: """Return AA agentic index score, or a low default for unknown models.""" - _load_aa_scores() mid = m.get("id", "") return _AA_SCORES_CACHE.get(mid, 25.0) @@ -1443,6 +1487,8 @@ async def _save_best_model_to_disk(best_model: dict) -> None: async def get_best_free_model() -> dict: """Fetches currently free models from OpenRouter, matches against agentic scores, and returns the highest.""" global free_model_cache + if not _AA_SCORES_LOADED: + await asyncio.to_thread(_load_aa_scores) now = time.time() # Check if cache is still valid @@ -3871,9 +3917,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 +3950,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." @@ -3904,6 +3965,7 @@ async def save_annotations(payload: AnnotationPayload): else: existing[k] = item.model_dump() await _atomic_write_json_async(str(ann_path), existing) + _annotations_cache.pop(str(ann_path), None) return JSONResponse({"status": "ok", "saved": len(data)}) except Exception as e: diff --git a/router/memory_mcp.py b/router/memory_mcp.py index f2187282..326ad66a 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,16 +45,21 @@ 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, safe="") + return f"{PREFIX}:v2:{scope}:{safe_category}::{ts}:{h}" def _parse_key(key: str): """Parse a structured key back into (scope, category, timestamp, hash).""" try: 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 "" + prefix = parts[0].split(":") # memory:{scope}:{category} or memory:v2:{scope}:{category} + if len(prefix) > 1 and prefix[1] == "v2": + scope = prefix[2] if len(prefix) > 2 else "" + category = urllib.parse.unquote(prefix[3]) if len(prefix) > 3 else "" + else: + scope = prefix[1] if len(prefix) > 1 else "" + category = 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 +74,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_agy_proxy.py b/router/tests/test_agy_proxy.py index 404059eb..8ca6ee3b 100644 --- a/router/tests/test_agy_proxy.py +++ b/router/tests/test_agy_proxy.py @@ -1,4 +1,5 @@ -from unittest.mock import patch, MagicMock +import pytest +from unittest.mock import patch, MagicMock, AsyncMock from router.agy_proxy import _wrap_response, _is_quota_exhausted def test_wrap_response_basic(): @@ -58,45 +59,51 @@ def test_wrap_response_long_strings(): assert result["usage"]["completion_tokens"] == 250 assert result["usage"]["total_tokens"] == 375 -def test_is_quota_exhausted_stderr_markers(): +@pytest.mark.asyncio +async def test_is_quota_exhausted_stderr_markers(): markers = ["RESOURCE_EXHAUSTED", "code 429", "quota reached", "rate limit"] for marker in markers: - assert _is_quota_exhausted(0, "", marker) is True - assert _is_quota_exhausted(1, "", f"Error: {marker}") is True + assert await _is_quota_exhausted(0, "", marker) is True + assert await _is_quota_exhausted(1, "", f"Error: {marker}") is True -def test_is_quota_exhausted_success(): - assert _is_quota_exhausted(0, "some valid response", "") is False +@pytest.mark.asyncio +async def test_is_quota_exhausted_success(): + assert await _is_quota_exhausted(0, "some valid response", "") is False -def test_is_quota_exhausted_other_error(): - assert _is_quota_exhausted(1, "", "some other random error") is False +@pytest.mark.asyncio +async def test_is_quota_exhausted_other_error(): + assert await _is_quota_exhausted(1, "", "some other random error") is False -@patch("router.agy_proxy.os.path.exists") -@patch("builtins.open") +@patch("aiofiles.open") @patch("router.agy_proxy.time.time") -def test_is_quota_exhausted_empty_reads_log(mock_time, mock_open, mock_exists): +@pytest.mark.asyncio +async def test_is_quota_exhausted_empty_reads_log(mock_time, mock_open): mock_time.return_value = 1000.0 - mock_exists.return_value = True - mock_file = MagicMock() - mock_file.readlines.return_value = ["line 1\n", "RESOURCE_EXHAUSTED info\n", "line 3\n"] - mock_open.return_value.__enter__.return_value = mock_file + mock_file = AsyncMock() + mock_file.seek = AsyncMock() + mock_file.tell = AsyncMock(return_value=100) + mock_file.read = AsyncMock(return_value=b"line 1\nRESOURCE_EXHAUSTED info\nline 3\n") + mock_open.return_value.__aenter__.return_value = mock_file with patch("router.agy_proxy._last_log_check", 0): - assert _is_quota_exhausted(0, "", "") is True + assert await _is_quota_exhausted(0, "", "") is True @patch("router.agy_proxy.time.time") -def test_is_quota_exhausted_empty_throttled(mock_time): +@pytest.mark.asyncio +async def test_is_quota_exhausted_empty_throttled(mock_time): # Set time diff to be < 2.0 mock_time.return_value = 1001.0 with patch("router.agy_proxy._last_log_check", 1000.0): # Even without reading log, falls back to True - assert _is_quota_exhausted(0, "", "") is True + assert await _is_quota_exhausted(0, "", "") is True -@patch("router.agy_proxy.os.path.exists") +@patch("aiofiles.open") @patch("router.agy_proxy.time.time") -def test_is_quota_exhausted_empty_no_log_fallback(mock_time, mock_exists): +@pytest.mark.asyncio +async def test_is_quota_exhausted_empty_no_log_fallback(mock_time, mock_open): mock_time.return_value = 1000.0 - mock_exists.return_value = False + mock_open.side_effect = FileNotFoundError() with patch("router.agy_proxy._last_log_check", 0): - assert _is_quota_exhausted(0, "", "") is True + assert await _is_quota_exhausted(0, "", "") is True diff --git a/router/tests/test_detect_active_tool.py b/router/tests/test_detect_active_tool.py new file mode 100644 index 00000000..c5b359d1 --- /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 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..10070ed5 --- /dev/null +++ b/router/tests/test_get_redis.py @@ -0,0 +1,137 @@ +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(monkeypatch): + """Reset the global variables before and after each test.""" + monkeypatch.delenv("VALKEY_URL", raising=False) + 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..ea0f9ad2 --- /dev/null +++ b/router/tests/test_load_persisted_stats.py @@ -0,0 +1,64 @@ +import os +os.environ.setdefault("CONFIG_PATH", "/config/router_dir/config.yaml") + +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..236853eb --- /dev/null +++ b/router/tests/test_memory_mcp.py @@ -0,0 +1,341 @@ +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, + _is_memory_key, + _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}:v2:{scope}:{category}::{ts}:{h}" + assert key.startswith(f"{PREFIX}:v2:{SCOPE_GLOBAL}:{category}::") + + # Extract timestamp and hash part + match = re.match(rf"^{PREFIX}:v2:{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}:v2:{SCOPE_LOCAL}:{category}::") + + match = re.match(rf"^{PREFIX}:v2:{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}:v2:{SCOPE_GLOBAL}:cat1::1620000000123:5e5dad075ca7764bc51f" + + key2 = _make_key("cat2", False, "data") + assert key2 == f"{PREFIX}:v2:{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 + +def test_is_memory_key_types(): + """Test _is_memory_key works with both string and non-string inputs.""" + assert _is_memory_key("memory:local:test") is True + assert _is_memory_key("other:prefix") is False + assert _is_memory_key(None) is False + assert _is_memory_key(12345) is False + assert _is_memory_key([]) is False + + +@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"}, + ), + ( + "memory:v2:local:proj%3Aalpha%2F100%25%20ready::20240101T120000Z:abc123hash", + {"scope": "local", "category": "proj:alpha/100% ready", "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", + "v2_escaped_category", + ] +) +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_fallback(): + 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_fallback(): + 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" + + +def test_make_key_and_parse_key_round_trip(): + """Verify that _make_key and _parse_key correctly quote and unquote complex categories.""" + category = "proj:alpha/100% ready" + key = _make_key(category, is_global=False, data="test-data") + + # Assert that the category in the key is URL-encoded + assert "proj%3Aalpha%2F100%25%20ready" in key + + # Assert that the parsed key returns the original unencoded category + parsed = _parse_key(key) + assert parsed["scope"] == "local" + assert parsed["category"] == category 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_classifier.py b/scripts/benchmark_classifier.py index cb1f5a4a..8000fa5e 100644 --- a/scripts/benchmark_classifier.py +++ b/scripts/benchmark_classifier.py @@ -1,6 +1,8 @@ """Benchmark gemma4-26a4b-routing classifier against labeled dataset.""" import os import json, urllib.request, time, sys +import concurrent.futures +import threading from collections import defaultdict, Counter from pathlib import Path @@ -54,41 +56,57 @@ def classify(prompt): per_tier = {t: {"correct": 0, "total": 0} for t in TIERS} confusion = defaultdict(Counter) # confusion[expected][predicted] -for i, item in enumerate(dataset.get("prompts", [])): +def process_item(item): prompt = item["prompt"] - # Support both old schema ("tier") and new schema ("llm_tier" / "clf_tier") - expected = item.get("tier") or item.get("llm_tier") or item.get("clf_tier", "") + expected = item.get("tier") or item.get("clf_tier") or item.get("llm_tier", "") try: predicted = classify(prompt) except Exception as e: predicted = f"ERROR: {str(e)[:50]}" - results.append({ - "prompt": prompt[:100], - "expected": expected, - "predicted": predicted, - }) + return expected, predicted - # Only score against known tiers — skip ERROR/unknown labels gracefully - if expected not in per_tier: +results_list = [None] * total + +with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor: + sem = threading.Semaphore(5) + + def process_item_with_rate_limit(index_and_item): + i, item = index_and_item + with sem: + # Add a small delay between acquisitions to stagger server load + time.sleep(0.05) + expected, predicted = process_item(item) + return i, item, expected, predicted + + futures = [executor.submit(process_item_with_rate_limit, (i, item)) for i, item in enumerate(dataset.get("prompts", []))] + + completed_count = 0 + for future in concurrent.futures.as_completed(futures): + i, item, expected, predicted = future.result() + + results_list[i] = { + "prompt": item["prompt"][:100], + "expected": expected, + "predicted": predicted, + } + + if expected in per_tier: + per_tier[expected]["total"] += 1 + if predicted == expected: + correct += 1 + per_tier[expected]["correct"] += 1 confusion[expected][predicted] += 1 - continue - per_tier[expected]["total"] += 1 - if predicted == expected: - correct += 1 - per_tier[expected]["correct"] += 1 + completed_count += 1 - confusion[expected][predicted] += 1 - - # Progress - if (i + 1) % 20 == 0: - scored_so_far = sum(t["total"] for t in per_tier.values()) - acc = (correct / scored_so_far * 100) if scored_so_far > 0 else 0.0 - print(f" {i+1}/{total} — accuracy {acc:.1f}%") + if completed_count % 20 == 0: + scored_so_far = sum(t["total"] for t in per_tier.values()) + acc = (correct / scored_so_far * 100) if scored_so_far > 0 else 0.0 + print(f" {completed_count}/{total} — accuracy {acc:.1f}%") - time.sleep(0.05) # minimal rate-limit (model handles concurrency via llama-server slots) +results = [r for r in results_list if r is not None] # Report scored_total = sum(t["total"] for t in per_tier.values()) diff --git a/scripts/benchmark_tokens.py b/scripts/benchmark_tokens.py new file mode 100644 index 00000000..1fe8da63 --- /dev/null +++ b/scripts/benchmark_tokens.py @@ -0,0 +1,77 @@ +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 + + if not all_passed: + raise ValueError("Token estimation accuracy benchmark failed") + +if __name__ == "__main__": + try: + verify_accuracy() + sys.exit(0) + except ValueError 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..c64bac43 --- /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") or [] + + # 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..e69929ea 100755 --- a/start-stack.sh +++ b/start-stack.sh @@ -31,14 +31,13 @@ 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 [ -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" ] || [ -z "$LANGFUSE_INIT_USER_PASSWORD" ] || [ -z "$REDIS_AUTH" ] || [ -z "$CLICKHOUSE_PASSWORD" ] || [ -z "$LANGFUSE_PUBLIC_KEY" ] || [ -z "$LANGFUSE_SECRET_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 if [ -t 0 ]; then echo "🔑 OpenRouter API Key not found." @@ -101,17 +100,21 @@ 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 ! 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 - # Ensure the env file exists and has secure permissions (owner read/write only) touch "$ENV_FILE" chmod 600 "$ENV_FILE" +generate_uuid() { + local val + val=$(openssl rand -hex 16 2>/dev/null) + local status=$? + if [ $status -ne 0 ] || [ ${#val} -ne 32 ]; then + echo "❌ Error: Failed to generate secure random UUID (openssl rand returned exit status $status, length ${#val})." >&2 + return 1 + fi + echo "${val:0:8}-${val:8:4}-${val:12:4}-${val:16:4}-${val:20:12}" +} + if [ -z "$NEXTAUTH_SECRET" ]; then NEXTAUTH_SECRET="$(openssl rand -base64 32)" echo "NEXTAUTH_SECRET=\"$NEXTAUTH_SECRET\"" >> "$ENV_FILE" @@ -141,12 +144,85 @@ 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 + +if [ -z "$LANGFUSE_PUBLIC_KEY" ]; then + uuid=$(generate_uuid) + if [ $? -ne 0 ] || [ -z "$uuid" ]; then + echo "❌ Error: Failed to generate LANGFUSE_PUBLIC_KEY." >&2 + exit 1 + fi + LANGFUSE_PUBLIC_KEY="pk-lf-$uuid" + echo "LANGFUSE_PUBLIC_KEY=\"$LANGFUSE_PUBLIC_KEY\"" >> "$ENV_FILE" + chmod 600 "$ENV_FILE" + echo "✓ Generated new LANGFUSE_PUBLIC_KEY and saved to $ENV_FILE" +fi + +if [ -z "$LANGFUSE_SECRET_KEY" ]; then + uuid=$(generate_uuid) + if [ $? -ne 0 ] || [ -z "$uuid" ]; then + echo "❌ Error: Failed to generate LANGFUSE_SECRET_KEY." >&2 + exit 1 + fi + LANGFUSE_SECRET_KEY="sk-lf-$uuid" + echo "LANGFUSE_SECRET_KEY=\"$LANGFUSE_SECRET_KEY\"" >> "$ENV_FILE" + chmod 600 "$ENV_FILE" + echo "✓ Generated new LANGFUSE_SECRET_KEY and saved to $ENV_FILE" +fi + +if [ -z "$OLLAMA_API_KEY" ]; then + if [ -t 0 ]; then + echo "🔑 OLLAMA_API_KEY not found." + echo -n "Please enter your Ollama API Key (input will be hidden): " + read -rs OLLAMA_API_KEY + echo "" + echo "OLLAMA_API_KEY=\"$OLLAMA_API_KEY\"" >> "$ENV_FILE" + chmod 600 "$ENV_FILE" + echo "✓ Ollama API key saved securely to $ENV_FILE" + else + echo "❌ Error: OLLAMA_API_KEY is not set in your environment or in $ENV_FILE." + echo "Please run this script interactively first, or create the file manually:" + echo " echo 'OLLAMA_API_KEY=your_key_here' >> $ENV_FILE" + echo " chmod 600 $ENV_FILE" + exit 1 + fi +fi + + # DYNAMIC_LITELLM_MASTER_KEY_PLACEHOLDER in router config is resolved at runtime from env @@ -255,7 +331,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 +432,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 OLLAMA_API_KEY LANGFUSE_PUBLIC_KEY LANGFUSE_SECRET_KEY python3 - "$WORKDIR/pod.yaml" <<'PY' import os, sys, urllib.parse uid = os.getuid() @@ -366,28 +442,40 @@ placeholders = [ "/home/gpav/Vrac/LAB/AI/LLM-Routing", "/home/gpav/", "/run/user/1000", - "sk-lit...33bf", - "postgres:***", + "LITELLM_MASTER_KEY_PLACEHOLDER", + "POSTGRES_PASSWORD_RAW_PLACEHOLDER", + "POSTGRES_PASSWORD_ENCODED_PLACEHOLDER", "NEXTAUTH_SECRET_PLACEHOLDER", "SALT_PLACEHOLDER", "ENCRYPTION_KEY_PLACEHOLDER", - "postgres-password-***" + "OLLAMA_API_KEY_PLACEHOLDER", + "LANGFUSE_PUBLIC_KEY_PLACEHOLDER", + "LANGFUSE_SECRET_KEY_PLACEHOLDER", + "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"] + "/") text = text.replace("/run/user/1000", f"/run/user/{uid}") -text = text.replace("sk-lit...33bf", os.environ["LITELLM_MASTER_KEY"]) +text = text.replace("LITELLM_MASTER_KEY_PLACEHOLDER", os.environ["LITELLM_MASTER_KEY"]) +text = text.replace("POSTGRES_PASSWORD_RAW_PLACEHOLDER", os.environ["POSTGRES_PASSWORD"]) # URL-encode the postgres password for DSN insertion -encoded_password = urllib.parse.quote_plus(os.environ['POSTGRES_PASSWORD']) -text = text.replace("postgres:***", f"postgres:{encoded_password}") -text = text.replace("postgres-password-***", os.environ["POSTGRES_PASSWORD"]) +encoded_password = urllib.parse.quote(os.environ['POSTGRES_PASSWORD'], safe="") +text = text.replace("POSTGRES_PASSWORD_ENCODED_PLACEHOLDER", encoded_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("OLLAMA_API_KEY_PLACEHOLDER", os.environ["OLLAMA_API_KEY"]) +text = text.replace("LANGFUSE_PUBLIC_KEY_PLACEHOLDER", os.environ["LANGFUSE_PUBLIC_KEY"]) +text = text.replace("LANGFUSE_SECRET_KEY_PLACEHOLDER", os.environ["LANGFUSE_SECRET_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 92% rename from test_compute_free_model_score.py rename to tests/test_compute_free_model_score.py index 2fab732c..1ea59751 100644 --- a/test_compute_free_model_score.py +++ b/tests/test_compute_free_model_score.py @@ -18,6 +18,7 @@ def test_compute_free_model_score_known_model(): """Test when the model id exists in the cache.""" mock_data = json.dumps({"scores": {"model-a": 85.5}}) with patch("builtins.open", mock_open(read_data=mock_data)): + router_main._load_aa_scores() score = compute_free_model_score({"id": "model-a"}) assert score == 85.5 @@ -25,6 +26,7 @@ def test_compute_free_model_score_unknown_model(): """Test when the model id is not in the cache.""" mock_data = json.dumps({"scores": {"model-a": 85.5}}) with patch("builtins.open", mock_open(read_data=mock_data)): + router_main._load_aa_scores() score = compute_free_model_score({"id": "model-b"}) assert score == 25.0 @@ -32,12 +34,14 @@ def test_compute_free_model_score_missing_id(): """Test when the model dictionary is missing an 'id'.""" mock_data = json.dumps({"scores": {"model-a": 85.5}}) with patch("builtins.open", mock_open(read_data=mock_data)): + router_main._load_aa_scores() score = compute_free_model_score({"name": "just a name"}) assert score == 25.0 def test_compute_free_model_score_file_not_found(): """Test fallback when the aa_scores.json file is missing or fails to load.""" with patch("builtins.open", side_effect=FileNotFoundError): + router_main._load_aa_scores() score = compute_free_model_score({"id": "model-a"}) assert score == 25.0 assert router_main._AA_SCORES_LOADED is True diff --git a/tests/test_get_pr_status.py b/tests/test_get_pr_status.py new file mode 100644 index 00000000..f5fcbf55 --- /dev/null +++ b/tests/test_get_pr_status.py @@ -0,0 +1,101 @@ +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 + + +@patch("scripts.get_pr_status.run_cmd") +def test_get_pr_status_null_checks(mock_run_cmd, capsys): + mock_data = { + "state": "OPEN", + "reviewDecision": "REVIEW_REQUIRED", + "statusCheckRollup": None + } + 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: REVIEW_REQUIRED" in captured.out + assert "Checks: 0/0 passed" in captured.out 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