diff --git a/.agents/AGENTS.md b/.agents/AGENTS.md index c7de7c06..16a25bb5 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 tracked files from moved directories (e.g., under the old `tests/` or `scripts/` paths) or re-create files at the root level. Resolve conflicts by directing all changes and file operations 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..daa2e840 --- /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 at free-model/roster call sites, specifically `sync_adaptive_router_roster()` and `get_best_free_model()`, rather than inside `compute_free_model_score()`. The fix was to move `_load_aa_scores()` behind `await asyncio.to_thread(...)` and remove the direct synchronous call from the scoring path. +* **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:** **181 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..ece7d2fd --- /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. +* **Verify Relocated Files**: While Git's rename detection can assist with directory/file moves during rebase, it is not foolproof. The agent must manually verify relocated files (such as renamed test paths) after rebasing to ensure all changes landed in their correct new locations rather than being lost or orphaned. + +### 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..0b7d3836 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 --ignore=tests/test_antigravity.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..2d729763 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`. @@ -790,7 +790,31 @@ The cooldown mechanism works as follows: For auto-routing modes, the Triage Router handles failures by silently falling back to the classified free tier cascade. For direct requests to `llm-routing-ollama`, the router returns `429` immediately during cooldown, allowing LiteLLM to skip this model group and cascade to `openrouter-auto`. The cooldown status is visible via the `/metrics` endpoint (`ollama_cooldown_active` and `ollama_cooldown_remaining_seconds` gauges). -## 10. Performance Benchmarks\n\nThrough our local benchmarks, the following performance characteristics have been achieved: +## 9d. Live Stack Tier Testing & Verification + +The repository includes an automated integration script to test the 5-tier intent routing pipeline on the live gateway stack: +* **Location**: [test_reasoning_tiers.py](scripts/verification/test_reasoning_tiers.py) + +This script sends five sequential chat completion requests (from simple to advanced prompt complexities) to the gateway's `llm-routing-auto-free` auto-triage route, verifying that: +1. The local classifier correctly categorizes and labels the prompt intent. +2. The gateway successfully routes the prompt to the mapped LiteLLM fallback model group or provider. +3. The responses are returned successfully with acceptable latency. + +### How to Run + +Ensure the container stack is deployed and healthy: +```bash +./start-stack.sh +``` + +Execute the verification script: +```bash +./scripts/verification/test_reasoning_tiers.py +``` + +## 10. Performance Benchmarks + +Through our local benchmarks, the following performance characteristics have been achieved: | Triage Evaluation Layer | Latency Footprint | Hardware Offload | Efficiency Ratio | | :--- | :---: | :---: | :---: | 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..192cbb1b 100644 --- a/litellm/entrypoint.py +++ b/litellm/entrypoint.py @@ -98,6 +98,30 @@ 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) + else: + dt = dt.astimezone(timezone.utc) + return dt.isoformat().replace("+00:00", "Z") + 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..1f8ed291 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: @@ -187,8 +187,8 @@ spec: - name: CLICKHOUSE_USER value: clickhouse - name: CLICKHOUSE_PASSWORD - value: clickhouse - image: docker.io/clickhouse/clickhouse-server:26.5.1 + value: CLICKHOUSE_PASSWORD_PLACEHOLDER + image: docker.io/clickhouse/clickhouse-server:26.6.1.1193 livenessProbe: exec: command: @@ -196,7 +196,7 @@ spec: - --user - clickhouse - --password - - clickhouse + - CLICKHOUSE_PASSWORD_PLACEHOLDER - --query - SELECT 1 initialDelaySeconds: 20 @@ -222,7 +222,7 @@ spec: - --port - '6380' - --requirepass - - langfuse-redis-2026 + - REDIS_AUTH_PLACEHOLDER - --maxmemory-policy - noeviction - --loglevel @@ -242,7 +242,7 @@ spec: - -p - "6380" - -a - - langfuse-redis-2026 + - REDIS_AUTH_PLACEHOLDER - ping initialDelaySeconds: 2 periodSeconds: 5 @@ -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 @@ -287,7 +287,7 @@ spec: - name: CLICKHOUSE_USER value: clickhouse - name: CLICKHOUSE_PASSWORD - value: clickhouse + value: CLICKHOUSE_PASSWORD_PLACEHOLDER - name: CLICKHOUSE_CLUSTER_ENABLED value: 'false' - name: REDIS_HOST @@ -295,7 +295,7 @@ spec: - name: REDIS_PORT value: '6380' - name: REDIS_AUTH - value: langfuse-redis-2026 + value: REDIS_AUTH_PLACEHOLDER - name: LANGFUSE_INIT_ORG_ID value: org-local-dev-id - name: LANGFUSE_INIT_ORG_NAME @@ -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,13 +337,13 @@ 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 value: clickhouse - name: CLICKHOUSE_PASSWORD - value: clickhouse + value: CLICKHOUSE_PASSWORD_PLACEHOLDER - name: CLICKHOUSE_CLUSTER_ENABLED value: 'false' - name: REDIS_HOST @@ -351,7 +351,7 @@ spec: - name: REDIS_PORT value: '6380' - name: REDIS_AUTH - value: langfuse-redis-2026 + value: REDIS_AUTH_PLACEHOLDER - name: HOSTNAME value: 0.0.0.0 - name: PORT @@ -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..baf0ac91 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,19 @@ 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 + return False except Exception: pass # Empty stdout+stderr with rc=0 strongly suggests quota exhaustion @@ -346,8 +355,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 +430,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..4074efdc 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,63 @@ 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) + + +METADATA_OVERHEAD = 50 + + 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) + METADATA_OVERHEAD) async def sync_cooldowns_from_valkey() -> None: @@ -487,6 +533,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 = {} @@ -935,9 +983,12 @@ async def classify_request( try: client = get_http_client() + # Truncate prompt to prevent 400 Bad Request in llama-server due to cache-ram limit + max_chars = int(os.getenv("CLASSIFIER_INPUT_MAX_CHARS", "300")) + truncated_prompt = prompt[:max_chars] if len(prompt) > max_chars else prompt payload = { "model": router_model_name, - "messages": [{"role": "user", "content": system_prompt + prompt}], + "messages": [{"role": "user", "content": system_prompt + truncated_prompt}], "temperature": 0.0, "max_tokens": 15, } @@ -1407,7 +1458,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 +1493,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 @@ -2122,6 +2174,8 @@ async def agy_stream_generator(): pass logger.error(f"agy proxy failed: {type(e).__name__}, falling back to LiteLLM") + if target_model == "llm-routing-agy": + target_model = "agent-advanced-core" original_target_model = target_model # --- OLLAMA (via LiteLLM) --- @@ -3871,9 +3925,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 +3958,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 +3973,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..a438ad1d 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}:") @@ -80,7 +88,17 @@ def _memory_value(data: str, tags: list | None) -> str: def _parse_memory_value(raw: str) -> dict: """Decode stored value back into {data, tags}.""" try: - return json.loads(raw) + val = json.loads(raw) + if not isinstance(val, dict): + return {"data": val, "tags": []} + if "data" not in val: + val["data"] = str(val) + else: + if val["data"] is not None: + val["data"] = str(val["data"]) + if "tags" not in val or not isinstance(val["tags"], list): + val["tags"] = [] + return val except (json.JSONDecodeError, TypeError): return {"data": raw, "tags": []} @@ -219,14 +237,22 @@ async def handle_remove_memory_category(args: dict) -> str: scope_label = "global" if is_global else "local" return f"No memories found to remove in category '{category}' ({scope_label})." + deleted_count = 0 async with httpx.AsyncClient(timeout=30.0) as client: for entry in to_delete: key = entry["key"] - await client.delete(f"{API_URL}/{key}", timeout=5.0) + quoted_key = urllib.parse.quote(key, safe="") + r = await client.delete(f"{API_URL}/{quoted_key}", timeout=5.0) + if r.status_code == 200: + deleted_count += 1 + else: + scope_label = "global" if is_global else "local" + cat_label = f"category '{category}'" if category != "*" else "all categories" + return f"Error removing memory (deleted {deleted_count} of {len(to_delete)} from {cat_label} ({scope_label})): {r.text}" scope_label = "global" if is_global else "local" cat_label = f"category '{category}'" if category != "*" else "all categories" - return f"Removed {len(to_delete)} memory(ies) from {cat_label} ({scope_label})." + return f"Removed {deleted_count} memory(ies) from {cat_label} ({scope_label})." async def handle_remove_specific_memory(args: dict) -> str: @@ -261,7 +287,8 @@ async def handle_remove_specific_memory(args: dict) -> str: ) async with httpx.AsyncClient(timeout=10.0) as client: - r = await client.delete(f"{API_URL}/{target['key']}", timeout=5.0) + quoted_key = urllib.parse.quote(target['key'], safe="") + r = await client.delete(f"{API_URL}/{quoted_key}", timeout=5.0) if r.status_code == 200: return f"Removed memory in category '{category}' ({target['data'][:60]}...)." else: 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..684a013f --- /dev/null +++ b/router/tests/test_get_gemini_oauth_status.py @@ -0,0 +1,101 @@ +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) + +os.environ["CONFIG_PATH"] = os.path.join(router_path, "config.yaml") + +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..5e8e54d2 --- /dev/null +++ b/router/tests/test_load_persisted_stats.py @@ -0,0 +1,71 @@ +import os +import sys + +# 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) + +os.environ["CONFIG_PATH"] = os.path.join(router_path, "config.yaml") + +import json +import pytest +from unittest.mock import patch, mock_open + +import main +from 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(main.stats, clean_stats, clear=True): + yield main.stats + +def test_load_persisted_stats_file_not_exists(mock_stats): + with patch("main.os.path.exists", return_value=False) as mock_exists: + load_persisted_stats() + mock_exists.assert_called_once_with(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("main.os.path.exists", return_value=True): + with patch("main.open", mock_open(read_data=mock_json)): + with patch("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("main.os.path.exists", return_value=True): + with patch("main.open", side_effect=Exception("Mock read error")): + with patch("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..ff680ad4 --- /dev/null +++ b/router/tests/test_memory_mcp.py @@ -0,0 +1,449 @@ +import json +import re +import time +import urllib.parse +from unittest.mock import AsyncMock, MagicMock, patch +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 == {"data": "just a string", "tags": []} + + +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 + + +@pytest.mark.asyncio +async def test_handle_remove_memory_category_url_encoding(): + from memory_mcp import handle_remove_memory_category + + category = "test:cat/100%_done" + key = _make_key(category, is_global=False, data="some-value") + + mock_list_response = MagicMock() + mock_list_response.status_code = 200 + mock_list_response.json.return_value = { + "memories": [ + { + "key": key, + "value": _memory_value("some-value", ["tag1"]) + } + ] + } + + mock_delete_response = MagicMock() + mock_delete_response.status_code = 200 + + mock_client = AsyncMock() + mock_client.get.return_value = mock_list_response + mock_client.delete.return_value = mock_delete_response + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client_class.return_value.__aenter__.return_value = mock_client + + result = await handle_remove_memory_category({"category": category, "is_global": False}) + + assert "Removed 1 memory" in result + + expected_quoted_key = urllib.parse.quote(key, safe="") + mock_client.delete.assert_called_once() + called_url = mock_client.delete.call_args[0][0] + assert called_url.endswith(expected_quoted_key) + + +@pytest.mark.asyncio +async def test_handle_remove_specific_memory_url_encoding(): + from memory_mcp import handle_remove_specific_memory + + category = "test:cat/100%_done" + key = _make_key(category, is_global=False, data="some-value") + + mock_list_response = MagicMock() + mock_list_response.status_code = 200 + mock_list_response.json.return_value = { + "memories": [ + { + "key": key, + "value": _memory_value("some-value", ["tag1"]) + } + ] + } + + mock_delete_response = MagicMock() + mock_delete_response.status_code = 200 + + mock_client = AsyncMock() + mock_client.get.return_value = mock_list_response + mock_client.delete.return_value = mock_delete_response + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client_class.return_value.__aenter__.return_value = mock_client + + result = await handle_remove_specific_memory({ + "category": category, + "memory_content": "some-value", + "is_global": False + }) + + assert "Removed memory" in result + + expected_quoted_key = urllib.parse.quote(key, safe="") + mock_client.delete.assert_called_once() + called_url = mock_client.delete.call_args[0][0] + assert called_url.endswith(expected_quoted_key) + + +@pytest.mark.asyncio +async def test_handle_remove_memory_category_failure(): + from memory_mcp import handle_remove_memory_category + + key1 = _make_key("cat", is_global=False, data="val1") + key2 = _make_key("cat", is_global=False, data="val2") + + mock_list_response = MagicMock() + mock_list_response.status_code = 200 + mock_list_response.json.return_value = { + "memories": [ + {"key": key1, "value": _memory_value("val1", [])}, + {"key": key2, "value": _memory_value("val2", [])} + ] + } + + mock_response_200 = MagicMock() + mock_response_200.status_code = 200 + mock_response_500 = MagicMock() + mock_response_500.status_code = 500 + mock_response_500.text = "Internal Server Error" + + mock_client = AsyncMock() + mock_client.get.return_value = mock_list_response + mock_client.delete.side_effect = [mock_response_200, mock_response_500] + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client_class.return_value.__aenter__.return_value = mock_client + + result = await handle_remove_memory_category({"category": "cat", "is_global": False}) + + assert "Error removing memory" in result + assert "deleted 1 of 2" in result + assert "Internal Server Error" in result 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..4786a626 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,67 @@ 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", [])): - 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", "") - +def process_item(item): try: + if not isinstance(item, dict): + raise TypeError("Item is not a dictionary") + prompt = item["prompt"] + expected = item.get("tier") or item.get("clf_tier") or item.get("llm_tier", "") predicted = classify(prompt) except Exception as e: + expected = "" + if isinstance(item, dict): + expected = item.get("tier") or item.get("clf_tier") or item.get("llm_tier", "") predicted = f"ERROR: {str(e)[:50]}" - results.append({ - "prompt": prompt[:100], - "expected": expected, - "predicted": predicted, - }) + return expected, predicted + +results_list = [None] * total - # Only score against known tiers — skip ERROR/unknown labels gracefully - if expected not in per_tier: +with concurrent.futures.ThreadPoolExecutor(max_workers=5) 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() + + prompt_val = "" + if isinstance(item, dict): + prompt_val = str(item.get("prompt", ""))[:100] + else: + prompt_val = f"" + + results_list[i] = { + "prompt": prompt_val, + "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..73be4033 --- /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, METADATA_OVERHEAD + +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) - METADATA_OVERHEAD # 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/scripts/verification/test_reasoning_tiers.py b/scripts/verification/test_reasoning_tiers.py new file mode 100755 index 00000000..6f702300 --- /dev/null +++ b/scripts/verification/test_reasoning_tiers.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +""" +Verification script to test the 5-tier intent classification and routing pipeline +of the LLM-Routing gateway. + +This script sends varying prompt complexities to trigger classification across +all 5 triage levels and validates the gateway's routing responses. +""" + +import httpx +import json +import time + +gateway_url = "http://127.0.0.1:5000/v1/chat/completions" +headers = { + "Authorization": "Bearer gateway-pass", + "Content-Type": "application/json" +} + +prompts = [ + { + "name": "Simple (1/5)", + "prompt": "Write a one-line hello world in Python." + }, + { + "name": "Medium (2/5)", + "prompt": "Refactor this Python function to add a docstring and use a type hint: def sum(a, b): return a + b" + }, + { + "name": "Complex (3/5)", + "prompt": "Write a multi-file Python script implementing a complete ETL data pipeline that parses a complex nested JSON file of users and logs, validates the data with Pydantic, writes it to a PostgreSQL database using asyncpg, handles connection pool retries with backoff, and includes unit tests with pytest mocks." + }, + { + "name": "Reasoning (4/5)", + "prompt": "Compare the design trade-offs of using Valkey/Redis Sentinel versus Valkey/Redis Cluster for high availability in a high-throughput, low-latency microservice architecture. Analyze failure detection mechanisms, failover times, read/write scaling, client complexity, and network partition (split-brain) scenarios." + }, + { + "name": "Advanced (5/5)", + "prompt": "Design a highly secure, zero-trust microservice topology for a federated agent-routing system deployed across multiple Kubernetes clusters in different cloud regions. Incorporate mutual TLS via Istio service mesh, automated certificate rotation with cert-manager, OIDC/OAuth2 authentication via an external identity provider, centralized audit logging/distributed tracing using Langfuse and Clickhouse, regional failover routing based on latency, and DDoS protection." + } +] + +def run_tests(): + print("🚀 Starting 5-tier reasoning test queries...") + print("=" * 60) + for p in prompts: + print(f"\n📂 Sending {p['name']} Prompt:") + print(f" Prompt: {p['prompt'][:100]}...") + + payload = { + "model": "llm-routing-auto-free", + "messages": [ + {"role": "user", "content": p["prompt"]} + ], + "temperature": 0.0, + "max_tokens": 150 + } + + start_time = time.time() + try: + r = httpx.post(gateway_url, json=payload, headers=headers, timeout=120.0) + elapsed = time.time() - start_time + print(f" Status: {r.status_code}") + + if r.status_code == 200: + data = r.json() + model_used = data.get("model", "unknown") + try: + msg = data["choices"][0]["message"] + content = msg.get("content") + if content is not None: + content_clean = content.strip().replace("\n", " ") + else: + content_clean = "" + print(f" Routed Model: {model_used}") + print(f" Response Preview: {content_clean[:120]}...") + except Exception as parse_inner: + print(f" Failed to parse response choices: {parse_inner}") + print(f" Raw Data: {data}") + print(f" Latency: {elapsed:.2f}s") + else: + print(f" Error: {r.text}") + except Exception as e: + print(f" Request Exception: {e}") + +if __name__ == "__main__": + run_tests() diff --git a/verify_breaker.py b/scripts/verification/verify_breaker.py similarity index 91% rename from verify_breaker.py rename to scripts/verification/verify_breaker.py index daf41bef..6ed9e7c7 100644 --- a/verify_breaker.py +++ b/scripts/verification/verify_breaker.py @@ -3,7 +3,7 @@ import sys from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from router.circuit_breaker import get_breaker diff --git a/watch_quota.sh b/scripts/watch_quota.sh similarity index 92% rename from watch_quota.sh rename to scripts/watch_quota.sh index 8bf9d090..cb95debd 100755 --- a/watch_quota.sh +++ b/scripts/watch_quota.sh @@ -2,7 +2,8 @@ # Polling loop — checks quota every 30s and runs tests when reset # Log file to watch LOG_FILE="$HOME/.gemini/antigravity-cli/cli.log" -TEST_SCRIPT="test_quota_reset.sh" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TEST_SCRIPT="$SCRIPT_DIR/test_quota_reset.sh" POLL_INTERVAL=30 # seconds echo "=== Quota Reset Watcher ===" diff --git a/start-stack.sh b/start-stack.sh index c9b826dc..7195ee19 100755 --- a/start-stack.sh +++ b/start-stack.sh @@ -12,6 +12,43 @@ set -e WORKDIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "$WORKDIR" +show_help() { + echo "Usage:" + echo " ./start-stack.sh → Restart existing pod (fast, preserves logs)" + echo " ./start-stack.sh --replace → Stop, clean up zombie ports, and recreate/redeploy pod" + echo " ./start-stack.sh --full-rebuild → Rebuild router image and recreate/redeploy pod" + echo " ./start-stack.sh --help | -h → Show this help message and exit" +} + +escape_env_val() { + local val="$1" + val="${val//\\/\\\\}" + val="${val//\"/\\\"}" + echo "$val" +} + + +if [ $# -gt 1 ]; then + echo "❌ Error: Too many arguments supplied (expected at most 1, got $#)" + show_help + exit 1 +fi + +FULL_REBUILD=false +REPLACE_MODE=false +if [ "${1:-}" = "--full-rebuild" ]; then + FULL_REBUILD=true +elif [ "${1:-}" = "--replace" ]; then + REPLACE_MODE=true +elif [ "${1:-}" = "--help" ] || [ "${1:-}" = "-h" ]; then + show_help + exit 0 +elif [ -n "${1:-}" ]; then + echo "❌ Error: Unknown argument '${1}'" + show_help + exit 1 +fi + # Ensure local volume directories exist on the host for Podman mounts mkdir -p valkey-data postgres-data langfuse-data clickhouse-data redis-lf-data minio-data @@ -31,21 +68,21 @@ 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." echo -n "Please enter your OpenRouter API Key (input will be hidden): " read -rs OPENROUTER_API_KEY echo "" - echo "OPENROUTER_API_KEY=\"$OPENROUTER_API_KEY\"" >> "$ENV_FILE" + local escaped_key=$(escape_env_val "$OPENROUTER_API_KEY") + echo "OPENROUTER_API_KEY=\"$escaped_key\"" >> "$ENV_FILE" chmod 600 "$ENV_FILE" echo "✓ API key saved securely to $ENV_FILE" else @@ -74,7 +111,7 @@ if [ -f "$OAUTH_CREDS" ]; then fi fi if $NEED_SYNC; then - python3 sync_gemini_token.py || echo "⚠️ Warning: Failed to sync Gemini token from keyring" + python3 scripts/sync_gemini_token.py || echo "⚠️ Warning: Failed to sync Gemini token from keyring" fi ACTIVE_OAUTH="" @@ -101,17 +138,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,23 +182,103 @@ 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 -# DYNAMIC_LITELLM_MASTER_KEY_PLACEHOLDER in router config is resolved at runtime from env +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 -FULL_REBUILD=false -REPLACE_MODE=false -if [ "${1:-}" = "--full-rebuild" ]; then - FULL_REBUILD=true -elif [ "${1:-}" = "--replace" ]; then - REPLACE_MODE=true +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." + while [ -z "$OLLAMA_API_KEY" ]; do + echo -n "Please enter your Ollama API Key (input will be hidden): " + read -rs OLLAMA_API_KEY + echo "" + if [ -z "$OLLAMA_API_KEY" ]; then + echo "❌ Error: API key cannot be empty. Please try again." + fi + done + local escaped_key=$(escape_env_val "$OLLAMA_API_KEY") + echo "OLLAMA_API_KEY=\"$escaped_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 + +if [ -z "$CLASSIFIER_INPUT_MAX_CHARS" ]; then + CLASSIFIER_INPUT_MAX_CHARS="300" + echo "CLASSIFIER_INPUT_MAX_CHARS=\"$CLASSIFIER_INPUT_MAX_CHARS\"" >> "$ENV_FILE" + echo "✓ Set default CLASSIFIER_INPUT_MAX_CHARS=300 and saved to $ENV_FILE" fi + + + +# DYNAMIC_LITELLM_MASTER_KEY_PLACEHOLDER in router config is resolved at runtime from env + +# Arguments parsed at top of script + # ── Cleanup zombie host-network ports ── # Podman with host networking can leave stuck LISTEN sockets after SIGKILL. # This covers ALL ports used by the pod + cross-profile orphans from other @@ -255,7 +376,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,38 +477,58 @@ 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 CLASSIFIER_INPUT_MAX_CHARS REDIS_AUTH CLICKHOUSE_PASSWORD python3 - "$WORKDIR/pod.yaml" <<'PY' -import os, sys, urllib.parse +import os, sys, urllib.parse, json uid = os.getuid() with open(sys.argv[1], "r", encoding="utf-8") as f: text = f.read() + +def yaml_scalar(val): + return json.dumps(val) + 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", + "REDIS_AUTH_PLACEHOLDER", + "CLICKHOUSE_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", yaml_scalar(os.environ["LITELLM_MASTER_KEY"])) +text = text.replace("POSTGRES_PASSWORD_RAW_PLACEHOLDER", yaml_scalar(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"]) -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"]) +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", yaml_scalar(os.environ["NEXTAUTH_SECRET"])) +text = text.replace("SALT_PLACEHOLDER", yaml_scalar(os.environ["SALT"])) +text = text.replace("ENCRYPTION_KEY_PLACEHOLDER", yaml_scalar(os.environ["ENCRYPTION_KEY"])) +text = text.replace("OLLAMA_API_KEY_PLACEHOLDER", yaml_scalar(os.environ["OLLAMA_API_KEY"])) +text = text.replace("LANGFUSE_PUBLIC_KEY_PLACEHOLDER", yaml_scalar(os.environ["LANGFUSE_PUBLIC_KEY"])) +text = text.replace("LANGFUSE_SECRET_KEY_PLACEHOLDER", yaml_scalar(os.environ["LANGFUSE_SECRET_KEY"])) +text = text.replace("MINIO_USER_PLACEHOLDER", yaml_scalar(os.environ["MINIO_ROOT_USER"])) +text = text.replace("MINIO_PASSWORD_PLACEHOLDER", yaml_scalar(os.environ["MINIO_ROOT_PASSWORD"])) +text = text.replace("LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER", yaml_scalar(os.environ["LANGFUSE_INIT_USER_PASSWORD"])) +text = text.replace("REDIS_AUTH_PLACEHOLDER", yaml_scalar(os.environ["REDIS_AUTH"])) +text = text.replace("CLICKHOUSE_PASSWORD_PLACEHOLDER", yaml_scalar(os.environ["CLICKHOUSE_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/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..a4c78758 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,19 @@ +import sys +import os +from pathlib import Path + +# Dynamic project root discovery +root = Path(__file__).resolve() +while root.parent != root and not (root / ".git").exists(): + root = root.parent + +if not (root / ".git").exists(): + raise RuntimeError("Could not find project root: .git directory not found in parent paths.") + +# Insert paths to sys.path +for path in [str(root), str(root / "router"), str(root / "scripts")]: + if path not in sys.path: + sys.path.insert(0, path) + +# Set common environment variables +os.environ.setdefault("CONFIG_PATH", str(root / "router" / "config.yaml")) diff --git a/test_a2_verify.py b/tests/test_a2_verify.py similarity index 83% rename from test_a2_verify.py rename to tests/test_a2_verify.py index 42cde1e1..95c6be31 100644 --- a/test_a2_verify.py +++ b/tests/test_a2_verify.py @@ -1,9 +1,5 @@ #!/usr/bin/env python3 """Verify circuit breaker integration into agy_proxy.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent / 'router')) - from circuit_breaker import get_breaker from agy_proxy import try_agy_proxy import asyncio, time 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 92% rename from test_antigravity.py rename to tests/test_antigravity.py index 8e0bc4b6..6b50d066 100644 --- a/test_antigravity.py +++ b/tests/test_antigravity.py @@ -4,6 +4,10 @@ import time def test_antigravity_connection(): + if os.environ.get("GITHUB_ACTIONS") == "true": + import pytest + pytest.skip("Skipping antigravity connection test in CI.") + creds_path = os.path.expanduser("~/.gemini/oauth_creds.json") if not os.path.exists(creds_path): print(f"Error: {creds_path} not found.") 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 99% rename from test_circuit_breaker.py rename to tests/test_circuit_breaker.py index 8cdcad60..695c0006 100644 --- a/test_circuit_breaker.py +++ b/tests/test_circuit_breaker.py @@ -18,8 +18,6 @@ import asyncio import pytest from unittest.mock import patch, AsyncMock -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent)) from router.circuit_breaker import get_breaker, TIER_COOLDOWNS, MAX_TIER 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..86600be9 --- /dev/null +++ b/tests/test_host_agy_daemon.py @@ -0,0 +1,431 @@ +import asyncio +import json +import os +import socket +import threading +import urllib.error +import urllib.request +from unittest.mock import AsyncMock + +import pytest + +import host_agy_daemon + +def make_run_request(daemon_server, payload): + return urllib.request.Request( + f"{daemon_server}/run", + data=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"}, + ) + +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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": False, "conversation_id": "conv_abc", "model_override": "gpt-4"}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": False, "timeout": 0.1}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": True, "model_override": "test-model"}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": True}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": True, "timeout": 0.1}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": False}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": True}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": True, "timeout": 0.1}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": True}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": False, "timeout": 0.1}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": False}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": False}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": False}) + + 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 = make_run_request(daemon_server, {"prompt": "test prompt", "stream": True, "conversation_id": "conv_789"}) + + 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 94% rename from test_models_proxy.py rename to tests/test_models_proxy.py index bfaf3b81..2ad68535 100644 --- a/test_models_proxy.py +++ b/tests/test_models_proxy.py @@ -4,13 +4,6 @@ from fastapi import Response from fastapi.responses import JSONResponse -# Set CONFIG_PATH for import -os.environ["CONFIG_PATH"] = "router/config.yaml" - -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent / "router")) - from main import get_http_client, proxy_models, HTTP_MAX_CONNECTIONS, HTTP_MAX_KEEPALIVE_CONNECTIONS, HTTP_KEEPALIVE_EXPIRY def test_http_client_limits(): 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..6a69cd26 --- /dev/null +++ b/tests/test_read_annotations_async.py @@ -0,0 +1,122 @@ +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 + +def make_mock_aiofiles_open(content: str): + mock_file = AsyncMock() + mock_file.read.return_value = content + mock_context_manager = MagicMock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_file) + mock_context_manager.__aexit__ = AsyncMock(return_value=False) + return MagicMock(return_value=mock_context_manager) + +@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_aiofiles_open = make_mock_aiofiles_open('{"annotation1": "data1"}') + + 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_aiofiles_open = make_mock_aiofiles_open('{"annotation2": "data2"}') + + 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