Async timeout issue (#53)

There's an issue with the way we have implemented the async check, it
needs a timeout to exit (early)

---------

Signed-off-by: Ivan Pedrazas <ipedrazas@gmail.com>

Author: ipedrazas

runtime_hook.py

@@ -1,3 +1,4 @@
+# ruff: noqa
 # Runtime hook to ensure stdlib modules are available
 import pickletools
 import pickle

scripts/bump_version.py

@@ -6,11 +6,11 @@
 
 import argparse
 import os
+from pathlib import Path
 import re
 import shutil
 import subprocess
 import sys
-from pathlib import Path
 
 
 def find_uv():
@@ -42,19 +42,19 @@ def get_latest_tag():
     if result.returncode != 0 or not result.stdout.strip():
         return None
 
-    tags = result.stdout.strip().split('\n')
+    tags = result.stdout.strip().split("\n")
     for tag in tags:
         # Find first tag that matches semver pattern
-        if re.match(r'^v?\d+\.\d+\.\d+$', tag):
+        if re.match(r"^v?\d+\.\d+\.\d+$", tag):
             return tag
     return None
 
 
 def parse_version(version_str):
     """Parse version string into major, minor, patch."""
     # Remove 'v' prefix if present
-    version_str = version_str.lstrip('v')
-    match = re.match(r'^(\d+)\.(\d+)\.(\d+)$', version_str)
+    version_str = version_str.lstrip("v")
+    match = re.match(r"^(\d+)\.(\d+)\.(\d+)$", version_str)
     if not match:
         raise ValueError(f"Invalid version format: {version_str}")
     return tuple(map(int, match.groups()))
@@ -64,11 +64,11 @@ def bump_version(version, bump_type):
     """Bump version based on type (major, minor, patch)."""
     major, minor, patch = version
 
-    if bump_type == 'major':
+    if bump_type == "major":
         return (major + 1, 0, 0)
-    elif bump_type == 'minor':
+    elif bump_type == "minor":
         return (major, minor + 1, 0)
-    elif bump_type == 'patch':
+    elif bump_type == "patch":
         return (major, minor, patch + 1)
     else:
         raise ValueError(f"Invalid bump type: {bump_type}")
@@ -92,11 +92,7 @@ def update_pyproject_toml(new_version):
 
     # Update version line
     updated_content = re.sub(
-        r'^version\s*=\s*"[^"]+"',
-        f'version = "{new_version}"',
-        content,
-        count=1,
-        flags=re.MULTILINE
+        r'^version\s*=\s*"[^"]+"', f'version = "{new_version}"', content, count=1, flags=re.MULTILINE
     )
 
     if content == updated_content:
@@ -110,7 +106,7 @@ def update_pyproject_toml(new_version):
 def create_and_push_tag(tag_name, push=True):
     """Create git tag and optionally push to remote."""
     # Create the tag
-    result = run_command(f'git tag {tag_name}', check=False)
+    result = run_command(f"git tag {tag_name}", check=False)
     if result.returncode != 0:
         print(f"Error creating tag: {result.stderr}", file=sys.stderr)
         return False
@@ -119,57 +115,35 @@ def create_and_push_tag(tag_name, push=True):
 
     if push:
         # Push the tag to remote
-        result = run_command(f'git push origin {tag_name}', check=False)
+        result = run_command(f"git push origin {tag_name}", check=False)
         if result.returncode != 0:
             print(f"Error pushing tag: {result.stderr}", file=sys.stderr)
             # Delete the local tag if push failed
-            run_command(f'git tag -d {tag_name}', check=False)
+            run_command(f"git tag -d {tag_name}", check=False)
             return False
         print(f"Pushed tag to remote: {tag_name}")
 
     return True
 
 
 def main():
-    parser = argparse.ArgumentParser(
-        description="Bump version in git tags and pyproject.toml"
-    )
-    parser.add_argument(
-        '--major',
-        action='store_true',
-        help='Bump major version (x.0.0)'
-    )
-    parser.add_argument(
-        '--minor',
-        action='store_true',
-        help='Bump minor version (0.x.0)'
-    )
-    parser.add_argument(
-        '--patch',
-        action='store_true',
-        help='Bump patch version (0.0.x) - default'
-    )
-    parser.add_argument(
-        '--no-push',
-        action='store_true',
-        help='Do not push tag to remote'
-    )
-    parser.add_argument(
-        '--dry-run',
-        action='store_true',
-        help='Show what would be done without making changes'
-    )
+    parser = argparse.ArgumentParser(description="Bump version in git tags and pyproject.toml")
+    parser.add_argument("--major", action="store_true", help="Bump major version (x.0.0)")
+    parser.add_argument("--minor", action="store_true", help="Bump minor version (0.x.0)")
+    parser.add_argument("--patch", action="store_true", help="Bump patch version (0.0.x) - default")
+    parser.add_argument("--no-push", action="store_true", help="Do not push tag to remote")
+    parser.add_argument("--dry-run", action="store_true", help="Show what would be done without making changes")
 
     args = parser.parse_args()
 
     # Determine bump type
-    bump_type = 'patch'  # default
+    bump_type = "patch"  # default
     if args.major:
-        bump_type = 'major'
+        bump_type = "major"
     elif args.minor:
-        bump_type = 'minor'
+        bump_type = "minor"
     elif args.patch:
-        bump_type = 'patch'
+        bump_type = "patch"
 
     # Get latest tag
     latest_tag = get_latest_tag()
@@ -192,7 +166,7 @@ def main():
         print(f"\n[DRY RUN] Would update pyproject.toml to version: {new_version_str}")
         print(f"[DRY RUN] Would create tag: {new_tag}")
         if not args.no_push:
-            print(f"[DRY RUN] Would push tag to remote")
+            print("[DRY RUN] Would push tag to remote")
         return
 
     # Update pyproject.toml
@@ -203,13 +177,13 @@ def main():
     # Update uv lock file
     print("Updating uv.lock...")
     uv_cmd = find_uv()
-    result = run_command(f'{uv_cmd} lock', check=False)
+    result = run_command(f"{uv_cmd} lock", check=False)
     if result.returncode != 0:
         print(f"Error updating uv.lock: {result.stderr}", file=sys.stderr)
         sys.exit(1)
 
     # Stage the pyproject.toml and uv.lock changes
-    result = run_command('git add pyproject.toml uv.lock', check=False)
+    result = run_command("git add pyproject.toml uv.lock", check=False)
     if result.returncode != 0:
         print(f"Error staging files: {result.stderr}", file=sys.stderr)
         sys.exit(1)
@@ -221,15 +195,15 @@ def main():
         print(f"Error committing version change: {result.stderr}", file=sys.stderr)
         sys.exit(1)
 
-    print(f"Committed version change")
+    print("Committed version change")
 
     # Create and push tag
     if not create_and_push_tag(new_tag, push=not args.no_push):
         sys.exit(1)
 
     # Push the commit
     if not args.no_push:
-        result = run_command('git push', check=False)
+        result = run_command("git push", check=False)
         if result.returncode != 0:
             print(f"Error pushing commit: {result.stderr}", file=sys.stderr)
             sys.exit(1)
@@ -238,5 +212,5 @@ def main():
     print(f"\n✓ Successfully bumped version to {new_version_str} ({new_tag})")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

src/neptune_mcp/client.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Any
 
+import httpx
 from neptune_common import (
     GetLogsResponse,
     GetProjectResponse,
@@ -11,7 +12,6 @@
     QueryDatabaseRequest,
 )
 import requests
-import httpx
 
 from neptune_mcp.config import SETTINGS
 

src/neptune_mcp/mcp.py

@@ -186,6 +186,19 @@ def add_new_resource(kind: str) -> dict[str, Any]:
 response = client.get_secret_value(SecretId="<aws_id>")
 secret = response['SecretString']
 ```
+""",
+        }
+    elif kind == "Database":
+        return {
+            "description": "Managed database for your applications.",
+            "neptune_json_configuration": """
+To add a database to a project, add the following to 'resources' in 'neptune.json':
+```json
+{
+    "kind": "Database",
+    "name": "<database_name>"
+}
+```
 """,
         }
     else:
@@ -367,9 +380,9 @@ async def deploy_project(neptune_json_path: str) -> dict[str, Any]:
         ]
         login_process = await asyncio.create_subprocess_shell(
             " ".join(login_cmd),
-            stdin = asyncio.subprocess.PIPE,
-            stdout  = asyncio.subprocess.DEVNULL,
-            stderr  = asyncio.subprocess.STDOUT,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.DEVNULL,
+            stderr=asyncio.subprocess.STDOUT,
             cwd=project_dir,
         )
 
@@ -420,18 +433,20 @@ async def deploy_project(neptune_json_path: str) -> dict[str, Any]:
         }
 
     # while deployment.status is not "Deployed", poll every 2 seconds
-    start_time = time.time()
     timeout = 180  # 3 minutes
-    while deployment.status != "Deployed" and deployment.status != "Error":
-        if time.time() - start_time > timeout:
-            log.error(f"Deployment timed out after {timeout} seconds")
-            return {
-                "status": "error",
-                "message": f"deployment timed out after {timeout} seconds while waiting for status 'Deployed'",
-                "next_step": "check the deployment status with 'get_deployment_status' and investigate any issues",
-            }
-        await asyncio.sleep(2)
-        deployment = await client.get_deployment_async(project_request.name, revision=deployment.revision)
+    try:
+        async with asyncio.timeout(timeout):
+            while deployment.status != "Deployed" and deployment.status != "Error":
+                await asyncio.sleep(2)
+                async with asyncio.timeout(10):
+                    deployment = await client.get_deployment_async(project_request.name, revision=deployment.revision)
+    except TimeoutError:
+        log.error(f"Deployment timed out after {timeout} seconds")
+        return {
+            "status": "error",
+            "message": f"deployment timed out after {timeout} seconds while waiting for status 'Deployed'",
+            "next_step": "check the deployment status with 'get_deployment_status' and investigate any issues",
+        }
 
     if deployment.status != "Deployed":
         log.error(f"Deployment failed with status: {deployment.status}")

src/neptune_mcp/utils.py

@@ -1,7 +1,5 @@
 import asyncio
 from dataclasses import dataclass
-import subprocess
-from typing import List, Union
 
 
 @dataclass
@@ -15,9 +13,7 @@ def success(self) -> bool:
         return self.returncode == 0
 
 
-async def run_command(
-    command: str | bytes, cwd: str | None = None
-) -> CommandResult:
+async def run_command(command: str | bytes, cwd: str | None = None) -> CommandResult:
     """
     Run a command and capture stdout, stderr, and return code.
 
@@ -31,8 +27,8 @@ async def run_command(
     """
     proc = await asyncio.create_subprocess_shell(
         command,
-        stderr  = asyncio.subprocess.PIPE,
-        stdout  = asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+        stdout=asyncio.subprocess.PIPE,
         cwd=cwd,
     )