[agent] cleanup: Fixed TODO structures and reorganized scripts

backend/scripts/benchmark.py

@@ -5,18 +5,20 @@
 """
 
 import asyncio
-import logging
 import json
+import logging
 import os
-from typing import List, Dict, Any
+from typing import Any, Dict, List
+
 from dotenv import load_dotenv
 
 # Load env vars before importing evaluators or agent components
 load_dotenv()
 
 from agent.graph import graph
+
 try:
-    from tests.evaluators import eval_quality, eval_groundedness
+    from tests.evaluators import eval_groundedness, eval_quality
 except ImportError:
     # This might happen if running script directly without module context
     # But usually handled by running as `python -m scripts.benchmark`
@@ -41,7 +43,7 @@ def load_dataset(path: str) -> List[Dict[str, Any]]:
             return []
 
     try:
-        with open(path, "r", encoding="utf-8") as f:
+        with open(path, encoding="utf-8") as f:
             return json.load(f)
     except Exception as e:
         logger.error(f"Failed to load dataset: {e}")

backend/scripts/check_path.py

@@ -1,6 +1,7 @@
 
-import sys
 import os
+import sys
+
 print(sys.path)
 try:
     import agent

scripts/test_available_models.py → backend/scripts/test_available_models.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
-"""
-Test which Gemini models are accessible via the google-genai SDK.
+"""Test which Gemini models are accessible via the google-genai SDK.
 """
 
 import os
@@ -14,13 +13,18 @@
 from google import genai
 
 # Add backend/src to path to import models
-PROJECT_ROOT = Path(__file__).parent.parent
+PROJECT_ROOT = Path(__file__).parent.parent.parent.resolve()
 BACKEND_SRC = PROJECT_ROOT / "backend" / "src"
 if str(BACKEND_SRC) not in sys.path:
     sys.path.append(str(BACKEND_SRC))
 
 try:
-    from agent.models import GEMINI_FLASH, GEMINI_FLASH_LITE, GEMINI_PRO, _DEPRECATED_MODELS
+    from agent.models import (
+        _DEPRECATED_MODELS,
+        GEMINI_FLASH,
+        GEMINI_FLASH_LITE,
+        GEMINI_PRO,
+    )
 except ImportError:
     print("[ERROR] Could not import agent.models. Check backend/src path.")
     sys.exit(1)
@@ -54,7 +58,7 @@ def main():
 
     if env_path.exists():
         env_vars = {}
-        with open(env_path, 'r', encoding='utf-8') as f:
+        with open(env_path, encoding='utf-8') as f:
             for line in f:
                 line = line.strip()
                 if line and not line.startswith('#') and '=' in line:

scripts/update_models.py → backend/scripts/update_models.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
-"""
-Script to update Gemini model configurations across the project.
+"""Script to update Gemini model configurations across the project.
 Usage: python update_models.py [strategy]
 Strategies:
   - flash (default): Gemini 2.5 Flash for all components (Best price-performance)
@@ -9,8 +8,8 @@
   - balanced: Flash-Lite for queries, Flash for reflection, Pro for answers
 """
 
-import sys
 import re
+import sys
 from pathlib import Path
 
 # Configuration Strategies - Only Gemini 2.5 models (1.5 and 2.0 are deprecated/inaccessible)
@@ -69,7 +68,7 @@
 # Assuming script is run from project root via scripts/update_models.sh or python scripts/update_models.py
 # If run directly from scripts/, we need parent.
 # But standard usage is from root. However, let's make it robust.
-PROJECT_ROOT = Path(__file__).parent.parent
+PROJECT_ROOT = Path(__file__).parent.parent.parent.resolve()
 BACKEND_DIR = PROJECT_ROOT / "backend/src/agent"
 FRONTEND_FILE = PROJECT_ROOT / "frontend/src/hooks/useAgentState.ts"
 ENV_FILE = PROJECT_ROOT / ".env"

backend/scripts/visualize_agent_graph.py

@@ -1,6 +1,6 @@
 
-import sys
 import os
+import sys
 from pathlib import Path
 
 # Add the src directory to sys.path to allow imports

backend/scripts/visualize_dependencies.py

@@ -1,13 +1,14 @@
 import ast
 import os
 import sys
-import pkg_resources
-import matplotlib.pyplot as plt
-import scipy.cluster.hierarchy as sch
-import numpy as np
 from collections import defaultdict
 from pathlib import Path
 
+import matplotlib.pyplot as plt
+import numpy as np
+import pkg_resources
+import scipy.cluster.hierarchy as sch
+
 # Set up paths
 BACKEND_ROOT = Path(__file__).resolve().parent.parent
 SRC_ROOT = BACKEND_ROOT / "src"
@@ -37,7 +38,7 @@
 def get_third_party_imports(file_path):
     """Parses a python file and returns a set of third-party base modules imported."""
     try:
-        with open(file_path, "r", encoding="utf-8") as f:
+        with open(file_path, encoding="utf-8") as f:
             tree = ast.parse(f.read())
     except Exception as e:
         print(f"Skipping {file_path}: {e}")

backend/src/agent/graph.py

@@ -153,7 +153,7 @@ def reflection_router(state: OverallState) -> list[Send] | str:
 )
 builder.add_edge("denoising_refiner", END)
 
-# TODO(priority=High, complexity=Medium): [SOTA Deep Research] Graph Wiring
+# TODO(priority=High, complexity=Medium, owner=team): [SOTA Deep Research] Graph Wiring
 # Add conditional edges to route from 'reflection' or 'update_plan' to 'research_subgraph'.
 # research_subgraph results should then flow back into 'update_plan' or merge into the state.
 

backend/src/agent/mcp_config.py

@@ -47,26 +47,26 @@ def validate(settings: MCPSettings) -> None:
 
 # Fine-grained implementation guide for MCP Integration:
 #
-# TODO(priority=High, complexity=Low): [MCP:1] Define SSE client interface
+# TODO(priority=High, complexity=Low, owner=team): [MCP:1] Define SSE client interface
 # - Create abstract base class for MCP transport
 # - Define methods: connect(), disconnect(), send_message(), receive_stream()
 #
-# TODO(priority=High, complexity=Medium): [MCP:2] Implement SSE transport
+# TODO(priority=High, complexity=Medium, owner=team): [MCP:2] Implement SSE transport
 # - Use httpx or aiohttp for Server-Sent Events
 # - Handle reconnection with exponential backoff
 # - Parse SSE event format (event:, data:, id:)
 #
-# TODO(priority=Medium, complexity=Medium): [MCP:3] Connection pooling
+# TODO(priority=Medium, complexity=Medium, owner=team): [MCP:3] Connection pooling
 # - Maintain pool of persistent connections
 # - Implement health checks and automatic reconnection
 # - Thread-safe connection acquisition/release
 #
-# TODO(priority=Medium, complexity=Low): [MCP:4] Error recovery
+# TODO(priority=Medium, complexity=Low, owner=team): [MCP:4] Error recovery
 # - Catch and log transport errors
 # - Retry failed tool calls with backoff
 # - Return graceful fallback on persistent failure
 #
-# TODO(priority=Low, complexity=Low): [MCP:5] Metrics and observability
+# TODO(priority=Low, complexity=Low, owner=team): [MCP:5] Metrics and observability
 # - Track connection latency, success/failure rates
 # - Integrate with Langfuse spans
 class McpConnectionManager:
@@ -94,7 +94,7 @@ def get_persistence_tools(self) -> List:
         ]
 
     async def get_tools(self):
-        # TODO(priority=High, complexity=Medium): [MCP:6] Implement actual SSE tool discovery
+        # TODO(priority=High, complexity=Medium, owner=team): [MCP:6] Implement actual SSE tool discovery
         # - Connect to MCP endpoint from settings
         # - Fetch tool list via SSE stream
         # - Convert to LangChain StructuredTool format

backend/src/agent/nodes.py

@@ -1,11 +1,11 @@
-# TODO(priority=Low, complexity=Low): See docs/tasks/upstream_compatibility.md for future splitting of this file into _nodes.py (upstream) and nodes.py (evolved).
+# TODO(priority=Low, complexity=Low, owner=team): See docs/tasks/upstream_compatibility.md for future splitting of this file into _nodes.py (upstream) and nodes.py (evolved).
 #
-# TODO(priority=Medium, complexity=Medium): [SOTA Deep Research] Benchmarking
+# TODO(priority=Medium, complexity=Medium, owner=team): [SOTA Deep Research] Benchmarking
 # See docs/tasks/04_SOTA_DEEP_RESEARCH_TASKS.md
 # Subtask: MLE-bench Integration (Evaluate on Kaggle engineering tasks).
 # Subtask: DeepResearch-Bench Setup (Load tasks from muset-ai space).
 
-# TODO(priority=Medium, complexity=High): Investigate and integrate 'deepagents' patterns if applicable.
+# TODO(priority=Medium, complexity=High, owner=team): Investigate and integrate 'deepagents' patterns if applicable.
 # See docs/tasks/04_SOTA_DEEP_RESEARCH_TASKS.md
 # Subtask: Review 'deepagents' repo for relevant nodes (e.g. hierarchical planning).
 # Subtask: Adapt useful patterns to `backend/src/agent/nodes.py`.
@@ -151,7 +151,7 @@ def scoping_node(state: OverallState, config: RunnableConfig) -> OverallState:
     If yes -> Generates questions and sets status to 'active' (interrupt).
     If no -> Sets status to 'complete' (proceed).
 
-    TODO(priority=High, complexity=High): [SOTA Deep Research] Verify full alignment with Open Deep Research (Clarification Loop).
+    TODO(priority=High, complexity=High, owner=team): [SOTA Deep Research] Verify full alignment with Open Deep Research (Clarification Loop).
     See docs/tasks/04_SOTA_DEEP_RESEARCH_TASKS.md
     Subtask: Implement `scoping_node` logic: Analyze input query. If ambiguous, generate clarifying questions and interrupt graph.
     """
@@ -1035,25 +1035,25 @@ def flow_update(state: OverallState, config: RunnableConfig) -> OverallState:
 
     Fine-grained implementation guide:
 
-    TODO(priority=High, complexity=Low): [flow_update:1] Extract current task from state
+    TODO(priority=High, complexity=Low, owner=team): [flow_update:1] Extract current task from state
     - Read `current_task_idx` and `plan` from state
     - Get the task object being evaluated
 
-    TODO(priority=High, complexity=Medium): [flow_update:2] Analyze task completion
+    TODO(priority=High, complexity=Medium, owner=team): [flow_update:2] Analyze task completion
     - Compare task query against `web_research_result`
     - Use fuzzy matching or LLM to determine if task is adequately answered
     - Return completion_score (0.0-1.0)
 
-    TODO(priority=High, complexity=Medium): [flow_update:3] Identify knowledge gaps
+    TODO(priority=High, complexity=Medium, owner=team): [flow_update:3] Identify knowledge gaps
     - Parse research results for "unclear", "contradictory", or "insufficient" signals
     - Generate list of follow-up questions if gaps detected
 
-    TODO(priority=Medium, complexity=High): [flow_update:4] DAG expansion logic
+    TODO(priority=Medium, complexity=High, owner=team): [flow_update:4] DAG expansion logic
     - If gaps detected: Create new tasks and insert into plan
     - If task complete: Mark status='done' and increment current_task_idx
     - If no more tasks: Set research_complete=True
 
-    TODO(priority=Low, complexity=Low): [flow_update:5] Return updated state
+    TODO(priority=Low, complexity=Low, owner=team): [flow_update:5] Return updated state
     - Return dict with updated `plan`, `current_task_idx`, `research_complete`
 
     See docs/tasks/04_SOTA_DEEP_RESEARCH_TASKS.md
@@ -1157,7 +1157,7 @@ def content_reader(state: OverallState, config: RunnableConfig) -> OverallState:
         return {"evidence_bank": extracted_evidence}
 
 
-# TODO(priority=High, complexity=Medium): [SOTA Deep Research] Recursive Trigger
+# TODO(priority=High, complexity=Medium, owner=team): [SOTA Deep Research] Recursive Trigger
 # Implement logic in reflection or a new 'router' node to decide when to call 'research_subgraph'.
 # This should happen when a complex sub-topic is identified that requires its own full research loop.
 def research_subgraph(state: OverallState, config: RunnableConfig) -> OverallState:

backend/src/agent/rag.py

@@ -537,7 +537,7 @@ class Resource:
 def create_rag_tool(resources):
     """Legacy compatibility stub - returns None.
 
-    TODO(priority=Low, complexity=Medium): [rag:legacy] Replace stub with real implementation
+    TODO(priority=Low, complexity=Medium, owner=team): [rag:legacy] Replace stub with real implementation
     - Migrate callers to use DeepSearchRAG directly
     - Remove this function once all callers are updated
     - Update tests that mock this function

backend/src/agent/security.py

@@ -65,7 +65,7 @@ def _is_ip_in_trusted_proxies(ip: str) -> bool:
 
 def extract_client_ip_from_forwarded(
     forwarded: str,
-    trusted_proxy_count: int = TRUSTED_PROXY_COUNT,
+    trusted_proxy_count: int | None = None,
     fallback_ip: str | None = None,
 ) -> str | None:
     """Extract the real client IP from X-Forwarded-For header using trust-bound extraction.
@@ -89,6 +89,10 @@ def extract_client_ip_from_forwarded(
     Returns:
         The extracted client IP, or fallback_ip if no valid candidate found.
     """
+    if trusted_proxy_count is None:
+        # Look at the module variable, don't hardcode the original. Tests modify it.
+        trusted_proxy_count = globals().get("TRUSTED_PROXY_COUNT", 0)
+
     if not forwarded:
         return fallback_ip
 

backend/src/evaluation/deep_research_bench.py

@@ -1,31 +1,31 @@
 # Fine-grained implementation guide for DeepResearch-Bench Evaluation:
 #
-# TODO(priority=High, complexity=Low): [deep_bench:1] Dataset loader
+# TODO(priority=High, complexity=Low, owner=team): [deep_bench:1] Dataset loader
 # - Connect to muset-ai/DeepResearch-Bench on HuggingFace
 # - Implement load_deep_research_dataset() -> List[Task]
 # - Each Task: {id, query, gold_report, evaluation_criteria}
 #
-# TODO(priority=High, complexity=Medium): [deep_bench:2] Agent runner
+# TODO(priority=High, complexity=Medium, owner=team): [deep_bench:2] Agent runner
 # - Import graph from agent.graph
 # - Configure for full research mode (scoping -> planning -> research -> synthesis)
 # - Capture final report and all intermediate artifacts
 #
-# TODO(priority=Medium, complexity=High): [deep_bench:3] Report scorer
+# TODO(priority=Medium, complexity=High, owner=team): [deep_bench:3] Report scorer
 # - Compare generated report against gold_report
 # - Use metrics: ROUGE-L, BERTScore, factual accuracy (via NLI)
 # - Return composite score (0.0-1.0)
 #
-# TODO(priority=Medium, complexity=Medium): [deep_bench:4] Citation verifier
+# TODO(priority=Medium, complexity=Medium, owner=team): [deep_bench:4] Citation verifier
 # - Check that all claims are backed by sources
 # - Verify source URLs are valid and content matches claims
 # - Return citation_coverage score
 #
-# TODO(priority=Medium, complexity=Low): [deep_bench:5] Metrics aggregator
+# TODO(priority=Medium, complexity=Low, owner=team): [deep_bench:5] Metrics aggregator
 # - Aggregate scores across all tasks
 # - Compute mean, std, percentiles
 # - Track token usage and latency
 #
-# TODO(priority=Low, complexity=Low): [deep_bench:6] Report generator
+# TODO(priority=Low, complexity=Low, owner=team): [deep_bench:6] Report generator
 # - Output results to JSON and Markdown
 # - Generate comparison charts (if multiple runs)
 #
@@ -34,30 +34,30 @@
 
 def evaluate_deep_research():
     """Evaluates the agent on DeepResearch-Bench (muset-ai)."""
-    # TODO(priority=High, complexity=Low): [deep_bench:1] Load dataset
+    # TODO(priority=High, complexity=Low, owner=team): [deep_bench:1] Load dataset
     dataset = []  # load_deep_research_dataset()
 
-    # TODO(priority=High, complexity=Medium): [deep_bench:2] Run agent
+    # TODO(priority=High, complexity=Medium, owner=team): [deep_bench:2] Run agent
     results = []
     for task in dataset:
         # report = run_full_research(task.query)
         # results.append({"task_id": task.id, "report": report})
         _ = task  # placeholder until implementation is complete
 
-    # TODO(priority=Medium, complexity=High): [deep_bench:3] Score reports
+    # TODO(priority=Medium, complexity=High, owner=team): [deep_bench:3] Score reports
     scores = []
     # for result in results:
     #     score = score_report(result["report"], gold_report)
     #     scores.append(score)
 
-    # TODO(priority=Medium, complexity=Medium): [deep_bench:4] Verify citations
+    # TODO(priority=Medium, complexity=Medium, owner=team): [deep_bench:4] Verify citations
     # for result in results:
     #     citation_score = verify_citations(result["report"])
 
-    # TODO(priority=Medium, complexity=Low): [deep_bench:5] Aggregate
+    # TODO(priority=Medium, complexity=Low, owner=team): [deep_bench:5] Aggregate
     # mean_score = sum(scores) / len(scores) if scores else 0
 
-    # TODO(priority=Low, complexity=Low): [deep_bench:6] Report
+    # TODO(priority=Low, complexity=Low, owner=team): [deep_bench:6] Report
     print("DeepResearch-Bench evaluation not yet implemented")