Skip to content

Support: Add middleware API integration for local CLI#58

Merged
sawradip merged 7 commits into
mainfrom
rad/sup_middleware
Aug 23, 2025
Merged

Support: Add middleware API integration for local CLI#58
sawradip merged 7 commits into
mainfrom
rad/sup_middleware

Conversation

@RadeenXALNW
Copy link
Copy Markdown
Collaborator

@RadeenXALNW RadeenXALNW commented Aug 23, 2025

Add middleware API integration for local CLI

Description

This PR adds middleware synchronization support to the local CLI, enabling seamless integration between local agent execution and the RunAgent cloud platform.

Changes

  • Middleware Sync Integration: Local CLI now connects to middleware via API when sync is enabled
  • Dual Mode Support:
    • Local-only mode: Data saved to local database when middleware sync is disabled
    • Cloud sync mode: Data automatically synced to middleware database when proper RunAgent token is provided
  • Cross-platform compatibility: Both streaming and non-streaming operations work with Python and Rust SDKs
  • Frontend visibility: Synced data appears in the RunAgent dashboard with enhanced UI presentation

Testing Status

  • Python SDK: Streaming and non-streaming modes tested
  • Rust SDK: Streaming and non-streaming modes tested
  • TypeScript/JavaScript SDK: Pending testing
  • Go SDK: Pending testing

Usage

When middleware sync is enabled with a valid RunAgent token, all local agent executions will automatically sync to the cloud platform, providing a unified view of agent activities across local and cloud environments.

Summary by CodeRabbit

  • New Features
    • CLI: clearer auth setup with targeted errors; richer Middleware Sync Status and tests; new db logs command; improved serve replace workflow with safe port allocation; enhanced db status with capacity insights.
    • SDK/Client: added explicit auth validation and connection diagnostics.
    • Server: improved streaming with per‑chunk data and full invocation tracking; new invocation endpoints; endpoint discovery printed at startup.
  • Chores
    • Updated default base URL to port 8333; improved cache directory resolution/creation with writability check.
  • Tests
    • Updated example scripts to showcase streaming flows and refreshed agent IDs.

@coderabbitai
Copy link
Copy Markdown

coderabbitai Bot commented Aug 23, 2025

Walkthrough

Adds a new MessageType/Data variant across Rust/Python schemas. Overhauls CLI setup, local sync, serve flow, and introduces a logs command. Updates constants (ports, cache dir). Revamps SDK auth caching and validation. Reworks middleware sync APIs and REST client diagnostics. Enriches local server and WebSocket streaming with tracking. Adjusts templates and tests.

Changes

Cohort / File(s) Summary
Schema message types
runagent-rust/runagent/src/types/schema.rs, runagent/utils/schema.py
Adds Data/data variant to MessageType; serde/snakes case serialization maintained.
CLI commands overhaul
runagent/cli/commands.py
Setup/auth flow wrapped with targeted errors; local_sync flags reduced; richer middleware status/tests; serve add replace/capacity logic and sync checks; new db.logs subcommand; improved db status tables.
Constants and cache handling
runagent/constants.py
Base URL ports changed 8330→8333; LOCAL_CACHE_DIRECTORY gains env-based fallback and mkdir with writability check.
SDK config/auth caching
runagent/sdk/config.py
_test_authentication returns structured dict using /users/profile; setup raises AuthenticationError with message; cached auth_validated flag; new validate_authentication API; user_info simplified.
Middleware sync service
runagent/sdk/deployment/middleware_sync.py
Uses cached auth for enablement; payload shapes simplified (agent_id/local_execution_id); endpoints updated; sync_invocation_complete signature changed; removed MiddlewareSync alias and sync_agent_logs; uniform test_connection result.
REST client diagnostics
runagent/sdk/rest_client.py
Error parsing improved; base URL handling tightened; validate_api_connection returns structured status; new debug_connection; internal token conversion hook; session cleanup.
Local server execution/sync
runagent/sdk/server/local_server.py
Startup sync payload enriched; two-step middleware invocation flow aligned; safer serialization and completion; new endpoints for invocations/stats; endpoint discovery printing.
WebSocket streaming with tracking
runagent/sdk/server/socket_utils.py
Streaming-with-tracking path overhauled: local/middleware invocation start, per-chunk DATA messages with IDs/timestamps, completion/error sync; signatures updated; serialization helper expanded.
Socket client minor import
runagent/sdk/socket_client.py
Adds uuid import; no functional/public changes.
Template config adjustments
templates/agno/default/runagent.config.json
Removes extractor from agno_assistant; drops timestamp mapping from agno_stream extractor.
Python test scripts updates
test_scripts/python/*
Align agent_ids; shift to streaming examples in ag2 and langgraph_advanced (adds main_streaming); update prompts.
Rust test script updates
test_scripts/rust/test_langgraph/src/main.rs
Update agent_id and constraints; comment out streaming section entirely.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor User
  participant CLI as CLI (commands.py)
  participant SDK as SDKConfig
  participant REST as RestClient
  Note over CLI: Setup / local_sync status
  User->>CLI: run db setup / local-sync --status
  CLI->>SDK: validate_authentication()
  SDK->>REST: health + /users/profile
  REST-->>SDK: { success, api_connected, api_authenticated, user_info? }
  SDK-->>CLI: { success, authenticated, user_info, base_url }
  CLI-->>User: Print structured status (API key, base URL, auth)
Loading
sequenceDiagram
  autonumber
  participant Server as LocalServer
  participant Sync as MiddlewareSyncService
  participant MW as Middleware API
  participant DB as Local DB
  Note over Server: Non-streaming invocation
  Server->>DB: start_invocation(input, entrypoint_tag)
  Server->>Sync: sync_invocation_start({agent_id, local_execution_id, input_data,...})
  Sync->>MW: POST /local-agents/invocations
  MW-->>Sync: { id: execution_id }
  Sync-->>Server: execution_id
  Server->>Server: Run agent
  Server->>DB: complete_invocation(output/error, timing)
  Server->>Sync: sync_invocation_complete(execution_id, completion_data)
  Sync->>MW: PATCH /local-agents/invocations/{execution_id}
  MW-->>Sync: { success }
Loading
sequenceDiagram
  autonumber
  participant WS as WebSocket Client
  participant Sock as AgentWebSocketHandler
  participant DB as Local DB
  participant Sync as MiddlewareSyncService
  Note over Sock: Streaming with tracking
  WS->>Sock: Connect + start stream
  Sock->>DB: start_invocation(...)
  alt sync enabled
    Sock->>Sync: sync_invocation_start(...)
    Sync-->>Sock: middleware_invocation_id
  end
  loop For each chunk
    Sock-->>WS: DATA {id, content, ts}
  end
  alt success
    Sock->>DB: complete_invocation(final_output)
    Sock->>Sync: sync_invocation_complete(..., output_data)
    Sock-->>WS: STATUS stream_complete
  else error/disconnect
    Sock->>DB: complete_invocation(error)
    Sock->>Sync: sync_invocation_complete(..., error_detail)
    Sock-->>WS: STATUS stream_error
  end
Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120–150 minutes

Poem

In ports we hop from 30 to 33,
I thump my paws at auth set free.
New DATA crumbs across the trail,
Streamed in chunks with fluffy detail.
Logs now bloom, replacements sing—
A whiskered nod to everything.
Hop, ship, and let the middleware ring! 🐇✨

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 Ruff (0.12.2)
runagent/cli/commands.py

�[1;31mruff failed�[0m
�[1mCause:�[0m Failed to parse /pyproject.toml
�[1mCause:�[0m TOML parse error at line 162, column 18
|
162 | target-version = "0.1.19"
| ^^^^^^^^
unknown variant 0.1.19, expected one of py37, py38, py39, py310, py311, py312, py313, py314

runagent/constants.py

�[1;31mruff failed�[0m
�[1mCause:�[0m Failed to parse /pyproject.toml
�[1mCause:�[0m TOML parse error at line 162, column 18
|
162 | target-version = "0.1.19"
| ^^^^^^^^
unknown variant 0.1.19, expected one of py37, py38, py39, py310, py311, py312, py313, py314

runagent/utils/schema.py

�[1;31mruff failed�[0m
�[1mCause:�[0m Failed to parse /pyproject.toml
�[1mCause:�[0m TOML parse error at line 162, column 18
|
162 | target-version = "0.1.19"
| ^^^^^^^^
unknown variant 0.1.19, expected one of py37, py38, py39, py310, py311, py312, py313, py314

  • 9 others

Tip

🔌 Remote MCP (Model Context Protocol) integration is now available!

Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats.

✨ Finishing Touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch rad/sup_middleware

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbitai help to get the list of available commands.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Status, Documentation and Community

  • Visit our Status Page to check the current availability of CodeRabbit.
  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link
Copy Markdown

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (6)
runagent/sdk/socket_client.py (1)

21-27: Fix WebSocket URL construction (lstrip bug, wrong scheme, and duplicated /api prefix).

Current logic:

  • Uses str.lstrip("http://")/lstrip("https://"), which strips any of those characters, not the substring; this can mangle domains.
  • Forces ws:// even when the base URL is https:// (should be wss://).
  • Can duplicate the API prefix when Config.get_base_url() already contains a path like /api/v1.

This will cause intermittent connection failures and downgrades TLS to plaintext.

Apply this diff to robustly parse the base URL, pick ws vs wss, and avoid duplicate prefixes:

@@
-        if not base_socket_url:
-            base_url = Config.get_base_url()
-            base_url = base_url.lstrip("http://").lstrip("https://")
-            base_socket_url = f"ws://{base_url}"
-
-        self.base_socket_url = base_socket_url.rstrip("/") + api_prefix
+        if not base_socket_url:
+            http_url = Config.get_base_url()
+        else:
+            # Allow callers to give either http(s)://... or ws(s)://...
+            http_url = base_socket_url
+
+        parsed = urlparse(http_url)
+        # Infer ws/wss from http/https when the scheme isn't already ws(s)
+        if parsed.scheme in ("ws", "wss"):
+            ws_scheme = parsed.scheme
+        else:
+            ws_scheme = "wss" if parsed.scheme == "https" else "ws"
+
+        host = parsed.hostname or parsed.netloc
+        port = f":{parsed.port}" if parsed.port else ""
+
+        # Normalize API prefix and avoid duplication if base URL already includes it
+        prefix = (api_prefix or "/api/v1")
+        if not prefix.startswith("/"):
+            prefix = f"/{prefix}"
+
+        base_path = (parsed.path or "").rstrip("/")
+        if base_path.endswith(prefix):
+            socket_path = base_path
+        elif base_path:
+            socket_path = f"{base_path}{prefix}"
+        else:
+            socket_path = prefix
+
+        self.base_socket_url = f"{ws_scheme}://{host}{port}{socket_path}"

Add the missing import at the top:

@@
-import json
-import uuid
+import json
+import uuid
+from urllib.parse import urlparse
runagent/constants.py (1)

31-36: Guard against empty RUNAGENT_CACHE_DIR yielding '.'

If RUNAGENT_CACHE_DIR is set to an empty string, Path("") resolves to ".", unintentionally using the current working directory as the cache. Treat empty strings as unset.

Apply this diff:

-_cache_dir = os.environ.get(ENV_LOCAL_CACHE_DIRECTORY)
-LOCAL_CACHE_DIRECTORY = str(
-    Path(_cache_dir)
-    if _cache_dir is not None
-    else Path(LOCAL_CACHE_DIRECTORY_PATH).expanduser()
-)
+_cache_dir = os.environ.get(ENV_LOCAL_CACHE_DIRECTORY)
+LOCAL_CACHE_DIRECTORY = str(
+    Path(_cache_dir)
+    if _cache_dir
+    else Path(LOCAL_CACHE_DIRECTORY_PATH).expanduser()
+)
runagent/sdk/rest_client.py (1)

150-160: JSON vs form mismatch when data is empty dict

The truthiness check causes {} to go as form-encoded (data=...) instead of JSON, and GET/DELETE conversion logic is also truthiness-based. Use “is not None” to preserve intent.

Apply this diff:

-                json=data if data and not files else None,
-                data=None if data and not files else data,
+                json=data if (data is not None and not files) else None,
+                data=None if (data is not None and not files) else data,

And above:

-        if method.upper() in ["GET", "DELETE"] and data and not params:
+        if method.upper() in ["GET", "DELETE"] and (data is not None) and not params:
runagent/sdk/server/local_server.py (1)

1101-1111: Remove duplicate shutdown_logging method.

There are two identical shutdown_logging methods defined (lines 1089-1099 and 1101-1111). The second one should be removed.

-    def shutdown_logging(self):
-        """Properly shutdown logging and flush remaining logs"""
-        try:
-            if hasattr(self, 'db_handler'):
-                self.db_handler.close()
-            if hasattr(self, 'agent_logger'):
-                for handler in self.agent_logger.handlers[:]:
-                    handler.close()
-                    self.agent_logger.removeHandler(handler)
-        except Exception as e:
-            console.print(f"⚠️ [yellow]Error during logging shutdown: {e}[/yellow]")
runagent/cli/commands.py (1)

563-569: Fix undefined variable capacity - should be capacity_info.

On line 568, the code references capacity which is not defined. It should be capacity_info which was defined on line 565.

             # Show capacity info
             # capacity = sdk.get_local_capacity()
             capacity_info = sdk.db_service.get_database_capacity_info()
 
             console.print(
-                f"📊 Capacity: [cyan]{capacity.get('current_count', 1)}/5[/cyan] slots used"
+                f"📊 Capacity: [cyan]{capacity_info.get('current_count', 1)}/5[/cyan] slots used"
             )
runagent/sdk/server/socket_utils.py (1)

62-69: Fix signature inconsistency - parameter name should be agent_execution_streamer.

The method signature uses stream_runner but the docstring and usage context suggest it should be agent_execution_streamer for consistency with other methods.

     async def handle_agent_stream_with_tracking(
         self, 
         websocket: WebSocket, 
         agent_id: str, 
-        stream_runner, 
+        agent_execution_streamer, 
         entrypoint_tag: str,
         db_service
     ):

And update the usage on line 160:

-            async for chunk in stream_runner(*input_args, **input_kwargs):
+            async for chunk in agent_execution_streamer(*input_args, **input_kwargs):
🧹 Nitpick comments (25)
runagent/sdk/socket_client.py (4)

34-35: Attach Authorization header to WebSocket handshake when api_key is set.

Middleware/cloud-sync often requires auth; without sending the token, connections can be rejected or untracked.

Apply this diff to pass the bearer token for both async and sync clients:

@@
-        async with websockets.connect(uri) as websocket:
+        headers = {}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        async with websockets.connect(uri, extra_headers=headers) as websocket:
@@
-        with connect(uri) as websocket:
+        headers = {}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        with connect(uri, additional_headers=headers) as websocket:

Also applies to: 82-83


45-50: Generate a unique correlation ID for stream start (use the new uuid import).

Hard-coding "stream_start" makes tracing difficult across concurrent streams. Use a UUID.

@@
-            start_msg = SafeMessage(
-                id="stream_start",
+            start_msg = SafeMessage(
+                id=str(uuid.uuid4()),
                 type=MessageType.STATUS,
                 timestamp="",
                 data=request.dict()
             )
@@
-            start_msg = SafeMessage(
-                id="stream_start",
+            start_msg = SafeMessage(
+                id=str(uuid.uuid4()),
                 type=MessageType.STATUS,
                 timestamp="",
                 data=request.dict()
             )

Also applies to: 94-99


120-121: Normalize chunk shape across async and sync streaming.

Async yields content when present, sync yields the whole data dict. Align them to reduce consumer conditionals.

-                    # Yield the actual chunk data
-                    yield safe_msg.data
+                    # Yield the actual chunk data (prefer content if present)
+                    yield safe_msg.data.get("content", safe_msg.data)

5-8: Remove unused import.

json isn't used in this module. Keeping it can trigger lint warnings.

-import json
+# import json  # unused

If you prefer, delete it entirely.

test_scripts/python/client_test_agno.py (2)

19-23: Make agent_id configurable to avoid hard-coding environment-specific UUIDs.

Hard-coded IDs tend to rot across environments. Read from an env var with a sensible default.

-from runagent import RunAgentClient
+import os
+from runagent import RunAgentClient
@@
-ra = RunAgentClient(
-    agent_id="c7a08c39-9086-436b-b64e-399779f5a7e8",
+ra = RunAgentClient(
+    agent_id=os.getenv("RUNAGENT_AGENT_ID", "c7a08c39-9086-436b-b64e-399779f5a7e8"),
     entrypoint_tag="agno_stream",
     local=True
     )

25-29: Print robustly across different chunk shapes.

Depending on server version, chunks may be dicts with content or already strings.

-for chunk in ra.run(
-    "Benefits of a long drive"
-):
-    print(chunk['content'], end="")
+for chunk in ra.run("Benefits of a long drive"):
+    text = chunk.get("content") if isinstance(chunk, dict) else str(chunk)
+    print(text, end="")
test_scripts/rust/test_langgraph/src/main.rs (2)

3-3: Remove unused import.

futures::StreamExt is unused now that streaming is disabled; this will trigger warnings.

-use futures::StreamExt;
+// use futures::StreamExt; // unused while streaming section is disabled

46-48: Clarify logging: you're printing the whole response, not only keys.

Either print keys or adjust the label.

-    println!("✅ Non-streaming response received");
-    println!("📄 Response keys: {:?}", 
-        response);
+    println!("✅ Non-streaming response received");
+    println!("📄 Response: {:#?}", response);
test_scripts/python/client_test_langgraph_advanced.py (2)

31-31: Drop unused asyncio import.

It's not used in this script.

-from runagent import RunAgentClient
-import asyncio
+from runagent import RunAgentClient

37-46: Handle chunk shape variations in streaming output.

Mirror the resilient printing approach so this keeps working as message formats evolve.

-    # Streaming version
-    for chunk in ra_client.run({
+    # Streaming version
+    for chunk in ra_client.run({
         "query": "My fridge is not getting cold.",
         "num_solutions": 3,
         "constraints": [{"type": "budget", "value": 100, "priority": "high"}],
         "user_context": {"experience_level": "beginner"},
         "metadata": {"test": True}
     }):
-        print(chunk)
+        print(chunk["content"] if isinstance(chunk, dict) and "content" in chunk else chunk)
templates/agno/default/runagent.config.json (1)

24-26: Streaming extractor narrowed to content only

For streaming, only "content" is extracted now. With the new MessageType.DATA flowing through streaming paths, ensure the emitter consistently places user-visible text at $.content or that the extractor should instead reference $.data.content for DATA messages if that’s the actual shape.

Would you like me to align the extractor to the final streaming message schema (DATA vs FINAL_RESPONSE) once we confirm the payload shape?

test_scripts/python/client_test_ag2.py (1)

34-39: Make streaming sample robust to DATA/STATUS chunks and missing keys

Some streamed chunks may be MessageType.DATA or STATUS and may not include ‘sender’/‘content’. Avoid KeyError and present meaningful output.

Apply this diff:

 for chunk in ra.run(
     message="Man breathe oxygen.",
     max_turns=3
 ):
-    print(f"{chunk['sender']}: {chunk['content']}")
-    print("-------")
+    sender = chunk.get("sender") or chunk.get("type", "unknown")
+    # Handle possible DATA-shaped payloads
+    content = chunk.get("content")
+    if content is None and isinstance(chunk.get("data"), dict):
+        content = chunk["data"].get("content") or chunk["data"]
+    print(f"{sender}: {content}")
+    print("-------")
runagent/constants.py (1)

38-47: mkdir alone doesn’t guarantee writability

If the directory already exists but is read-only, mkdir won’t fail. Proactively test write/delete a temp file to ensure it’s usable and fail fast with a clear error.

Apply this diff:

 try:
     Path(LOCAL_CACHE_DIRECTORY).mkdir(parents=True, exist_ok=True)
-except OSError as e:
+    # Verify writability
+    testfile = Path(LOCAL_CACHE_DIRECTORY) / ".write_test"
+    testfile.touch(exist_ok=True)
+    testfile.unlink(missing_ok=True)
+except OSError as e:
     if os.getenv('DISABLE_TRY_CATCH'):
         raise
     raise RuntimeError(
         f"Cache directory {LOCAL_CACHE_DIRECTORY} is not writable please "
         f"provide a path that is writable using {ENV_LOCAL_CACHE_DIRECTORY} "
         "environment variable."
     ) from e
runagent/sdk/rest_client.py (4)

78-83: Always set User-Agent (even without API key)

User-Agent is currently added only when an API key exists. Add it unconditionally so diagnostics and middleware analytics don’t depend on auth presence.

Apply this diff:

-        if self.api_key:
-            # Support both JWT tokens and API keys
-            self.session.headers.update({
-                "Authorization": f"Bearer {self.api_key}",
-                "User-Agent": "RunAgent-CLI/1.0"
-            })
+        # Always set UA
+        self.session.headers.update({"User-Agent": "RunAgent-CLI/1.0"})
+        if self.api_key:
+            # Support both JWT tokens and API keys
+            self.session.headers.update({"Authorization": f"Bearer {self.api_key}"})

89-112: Improve middleware error parsing and 401 detection

Some backends return {"detail": [{"msg": "..."}]} or lowercase messages. Make the 401 check case-insensitive and support list-shaped “detail”.

Apply this diff:

-        try:
+        try:
             error_data = response.json()
             if isinstance(error_data, dict):
                 # Handle middleware-style errors
                 if "detail" in error_data:
-                    error_message = error_data["detail"]
+                    detail = error_data["detail"]
+                    if isinstance(detail, list) and detail and isinstance(detail[0], dict) and "msg" in detail[0]:
+                        error_message = detail[0]["msg"]
+                    else:
+                        error_message = str(detail)
                 elif "message" in error_data:
                     error_message = error_data["message"]
                 elif "error" in error_data:
                     error_message = error_data["error"]
@@
-        if response.status_code == 401:
-            if "Not authenticated" in error_message:
+        if response.status_code == 401:
+            if "not authenticated" in error_message.lower():
                 raise AuthenticationError("API key is invalid or expired", response.status_code, response)
             else:
                 raise AuthenticationError(error_message, response.status_code, response)

706-711: Optional: Pre-validate/normalize API key to JWT once in init

You added _get_jwt_token_from_api_key but never use it. Consider normalizing the token once at init so HttpHandler carries the right credential.

Apply this diff:

-        self.http = HttpHandler(
-            api_key=self.api_key,  # Use API key directly - middleware handles conversion
+        normalized_key = self._get_jwt_token_from_api_key() or self.api_key
+        self.http = HttpHandler(
+            api_key=normalized_key,
             base_url=self.base_url
         )

926-961: _get_jwt_token_from_api_key: naming and behavior

This method doesn’t actually return a JWT; it returns the API key on success or even on error (best-effort). That’s fine if middleware accepts API-key-as-bearer, but the name suggests conversion. Consider renaming to _normalize_api_key_for_auth or return the JWT if/when the backend supports it.

runagent/utils/schema.py (1)

153-153: New MessageType.DATA — aligned with streaming payloads

Addition looks correct and aligns with the middleware/SDK streaming paths. Ensure emitters set “type”: "data" and that clients treat DATA separately from RAW_DATA/STRUCTURED_DATA to avoid ambiguity.

Would you like me to add short docstrings/comments distinguishing DATA vs RAW_DATA vs STRUCTURED_DATA to prevent misuse?

runagent/sdk/config.py (3)

137-207: Add timeout parameter configuration for the API request.

The method _test_authentication makes an HTTP request with a hardcoded timeout of 10 seconds. Consider making this configurable or at least consistent with other timeout values in the system.

-    def _test_authentication(self) -> t.Dict[str, t.Any]:
+    def _test_authentication(self, timeout: int = 10) -> t.Dict[str, t.Any]:
         """Test authentication with current configuration - SIMPLIFIED"""
         try:
             from .rest_client import RestClient
 
             client = RestClient(
                 api_key=self._config.get("api_key"),
                 base_url=self._config.get("base_url"),
             )
             
             # Test connection using the profile endpoint
-            response = client.http.get("/users/profile", timeout=10)
+            response = client.http.get("/users/profile", timeout=timeout)

176-186: Use consistent error handling for JSON parsing.

The error handling for JSON parsing on line 178 uses a bare except clause. Consider catching specific exceptions for better error handling.

             elif response.status_code == 401:
                 try:
                     error_data = response.json()
                     error_detail = error_data.get("detail", "Invalid API key")
-                except:
+                except (ValueError, TypeError, AttributeError):
                     error_detail = "API key authentication failed"

251-277: Consider caching authentication results in validate_authentication.

The validate_authentication method always calls _test_authentication, which makes a network request. Since this is a public method that might be called multiple times, consider adding an optional parameter to use cached results when available.

-    def validate_authentication(self) -> t.Dict[str, t.Any]:
-        """Public method to validate authentication - ONLY WHEN EXPLICITLY CALLED"""
+    def validate_authentication(self, force_refresh: bool = False) -> t.Dict[str, t.Any]:
+        """Public method to validate authentication - ONLY WHEN EXPLICITLY CALLED
+        
+        Args:
+            force_refresh: If True, always make a new authentication request.
+                          If False, use cached result if available within the last 5 minutes.
+        """
         if not self.is_configured():
             return {
                 "success": False,
                 "error": "No API key configured",
                 "configured": False
             }
         
+        # Use cached result if available and not forcing refresh
+        if not force_refresh and self._config.get("auth_validated"):
+            # Check if we validated recently (optional: add timestamp check)
+            return {
+                "success": True,
+                "authenticated": True,
+                "user_info": self.user_info,
+                "base_url": self.base_url
+            }
+        
         # Only validate if explicitly requested
         auth_result = self._test_authentication()
runagent/sdk/deployment/middleware_sync.py (2)

14-50: Consider extracting authentication check to a separate method.

The initialization logic is complex with nested conditions. Consider extracting the authentication validation logic into a separate method for better readability and maintainability.

 class MiddlewareSyncService:
     """Middleware sync service - UPDATED for simplified ID structure"""
     
     def __init__(self, config):
         self.config = config
         self.rest_client = None
+        self._initialize_sync_service()
+    
+    def _initialize_sync_service(self):
+        """Initialize the sync service based on configuration."""
         # Check if we have a valid API key
         self.api_key = getattr(self.config, 'api_key', None)
         if not self.api_key:
             console.print("[dim]No API key configured - middleware sync disabled[/dim]")
             self.sync_enabled = False
             self.enabled = False
             return
         
         # Initialize RestClient if we have an API key
         try:
             self.rest_client = RestClient(
                 base_url=self.config.base_url,
                 api_key=self.api_key
             )
             
             # If authentication was successful during setup, enable sync
             if hasattr(self.config, '_config') and self.config._config.get("auth_validated"):
                 console.print("[dim]Using cached authentication - middleware sync enabled[/dim]")
                 self.sync_enabled = True
                 self.enabled = True
             else:
                 console.print("[dim]No cached authentication - middleware sync disabled[/dim]")
                 self.sync_enabled = False
                 self.enabled = False
                 
         except Exception as e:
             console.print(f"[red]Could not initialize middleware sync: {e}[/red]")
             self.sync_enabled = False
             self.enabled = False
             self.rest_client = None

190-224: Add request ID for better traceability in async requests.

The _make_async_request method handles async requests but doesn't include any request ID for debugging. Consider adding a request ID for better traceability.

     async def _make_async_request(self, method: str, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
         """Make async HTTP request to middleware"""
+        import uuid
+        request_id = str(uuid.uuid4())[:8]
+        
         if not getattr(self, 'rest_client', None) or not self.is_sync_enabled():
             return {"success": False, "error": "Middleware sync not available"}
             
         try:
+            console.print(f"[dim]Request {request_id}: {method} {endpoint}[/dim]")
             if method == "POST":
                 response = await asyncio.to_thread(
                     self.rest_client.http.post, 
                     endpoint, 
                     data=data, 
                     timeout=30
                 )
             elif method == "PUT":
                 response = await asyncio.to_thread(
                     self.rest_client.http.put, 
                     endpoint, 
                     data=data, 
                     timeout=30
                 )
             else:
                 raise ValueError(f"Unsupported method: {method}")
runagent/sdk/server/local_server.py (1)

98-141: Consider extracting agent data preparation to a separate method.

The _sync_agent_to_middleware_and_wait method has complex agent data preparation logic that could be extracted for better readability.

     async def _sync_agent_to_middleware_and_wait(self):
         """Sync agent to middleware and wait for completion - FIXED for simplified ID structure"""
         if not hasattr(self, 'middleware_sync') or not self.middleware_sync:
             console.print("[dim]Middleware sync not available[/dim]")
             return False
             
         if not self.middleware_sync.sync_enabled:
             console.print("[dim]Middleware sync disabled (no API key configured)[/dim]")
             return False
 
         try:
-            # Prepare agent data for sync - FIXED to use simplified structure
-            agent_data = {
-                "local_agent_id": self.agent_id,  # This becomes the main agent ID in middleware
-                "name": self.agent_name,
-                "framework": self.agent_framework,
-                "version": self.agent_version,
-                "path": str(self.agent_path),
-                "host": self.host,
-                "port": self.port,
-                "entrypoints": [ep.dict() for ep in self.agent_architecture.entrypoints],
-                "status": "running",
-                "sync_timestamp": datetime.utcnow().isoformat()
-            }
+            agent_data = self._prepare_agent_sync_data()
 
             console.print(f"[cyan]Syncing agent {self.agent_id} to middleware...[/cyan]")
             console.print(f"[dim]Agent data: {agent_data['name']} ({agent_data['framework']})[/dim]")
             
             sync_result = await self.middleware_sync.sync_agent_startup(self.agent_id, agent_data)
             
             if sync_result:
                 console.print(f"✅ [green]Agent successfully synced to middleware[/green]")
                 self.agent_synced_to_middleware = True
                 return True
             else:
                 console.print(f"[yellow]Agent sync failed - logs will be local only[/yellow]")
                 self.agent_synced_to_middleware = False
                 return False
 
         except Exception as e:
             console.print(f"❌ [red]Failed to sync agent to middleware: {e}[/red]")
             self.agent_synced_to_middleware = False
             return False
+
+    def _prepare_agent_sync_data(self) -> Dict[str, Any]:
+        """Prepare agent data for middleware sync."""
+        return {
+            "local_agent_id": self.agent_id,  # This becomes the main agent ID in middleware
+            "name": self.agent_name,
+            "framework": self.agent_framework,
+            "version": self.agent_version,
+            "path": str(self.agent_path),
+            "host": self.host,
+            "port": self.port,
+            "entrypoints": [ep.dict() for ep in self.agent_architecture.entrypoints],
+            "status": "running",
+            "sync_timestamp": datetime.utcnow().isoformat()
+        }
runagent/sdk/server/socket_utils.py (1)

163-176: Consider making the chunk storage limit configurable.

The hardcoded limit of 5000 chunks for storage might not be suitable for all use cases. Consider making this configurable.

+    # Add as class constant
+    MAX_STORED_CHUNKS = 5000
+
     # In the method:
                 # Store chunk for final output tracking (limit to prevent memory issues)
-                if chunk_count <= 5000:
+                if chunk_count <= self.MAX_STORED_CHUNKS:
                     try:
                         serializable_chunk = self._convert_to_serializable(chunk)
                         stream_output_data.append(serializable_chunk)
                     except Exception as e:
                         console.print(f"⚠️ [yellow]Warning: Could not serialize chunk {chunk_count}: {e}[/yellow]")
                         stream_output_data.append({
                             "chunk_number": chunk_count,
                             "serialization_error": str(e),
                             "chunk_type": str(type(chunk)),
                             "chunk_preview": str(chunk)
                         })
📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

  • MCP integration is disabled by default for public repositories
  • Jira integration is disabled by default for public repositories
  • Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between b13930e and 285d273.

📒 Files selected for processing (15)
  • runagent-rust/runagent/src/types/schema.rs (1 hunks)
  • runagent/cli/commands.py (3 hunks)
  • runagent/constants.py (3 hunks)
  • runagent/sdk/config.py (6 hunks)
  • runagent/sdk/deployment/middleware_sync.py (5 hunks)
  • runagent/sdk/rest_client.py (7 hunks)
  • runagent/sdk/server/local_server.py (8 hunks)
  • runagent/sdk/server/socket_utils.py (6 hunks)
  • runagent/sdk/socket_client.py (1 hunks)
  • runagent/utils/schema.py (1 hunks)
  • templates/agno/default/runagent.config.json (1 hunks)
  • test_scripts/python/client_test_ag2.py (2 hunks)
  • test_scripts/python/client_test_agno.py (2 hunks)
  • test_scripts/python/client_test_langgraph_advanced.py (1 hunks)
  • test_scripts/rust/test_langgraph/src/main.rs (3 hunks)
🧰 Additional context used
🧬 Code graph analysis (8)
test_scripts/python/client_test_langgraph_advanced.py (1)
runagent-rust/runagent/src/client/runagent_client.rs (1)
  • entrypoint_tag (279-281)
runagent/sdk/config.py (4)
runagent/sdk/rest_client.py (2)
  • get (197-199)
  • AuthenticationError (46-48)
runagent/sdk/exceptions.py (1)
  • AuthenticationError (17-20)
runagent/sdk/sdk.py (1)
  • is_configured (112-114)
runagent/utils/config.py (1)
  • is_configured (260-268)
test_scripts/python/client_test_agno.py (1)
runagent-rust/runagent/src/client/runagent_client.rs (1)
  • entrypoint_tag (279-281)
runagent/sdk/server/socket_utils.py (5)
runagent/utils/serializer.py (2)
  • deserialize_message (117-153)
  • serialize_message (78-115)
runagent/sdk/db.py (2)
  • start_invocation (441-487)
  • complete_invocation (489-538)
runagent/sdk/deployment/middleware_sync.py (3)
  • is_sync_enabled (52-54)
  • sync_invocation_start (92-121)
  • sync_invocation_complete (123-140)
runagent/utils/schema.py (2)
  • SafeMessage (159-176)
  • MessageType (147-156)
runagent/sdk/server/local_server.py (1)
  • _convert_to_serializable (700-745)
runagent/sdk/server/local_server.py (5)
runagent/sdk/rest_client.py (1)
  • run_agent (854-891)
runagent/utils/schema.py (1)
  • AgentRunRequest (107-112)
runagent/sdk/db.py (3)
  • start_invocation (441-487)
  • complete_invocation (489-538)
  • record_agent_run (1337-1381)
runagent/sdk/deployment/middleware_sync.py (3)
  • is_sync_enabled (52-54)
  • sync_invocation_start (92-121)
  • sync_invocation_complete (123-140)
runagent/sdk/server/socket_utils.py (1)
  • _convert_to_serializable (659-704)
runagent/sdk/rest_client.py (2)
runagent/sdk/exceptions.py (1)
  • AuthenticationError (17-20)
runagent/utils/config.py (3)
  • Config (17-320)
  • get_api_key (166-181)
  • get_base_url (131-150)
runagent/cli/commands.py (5)
runagent/sdk/config.py (4)
  • user_info (290-296)
  • api_key (280-282)
  • base_url (285-287)
  • is_configured (208-210)
runagent/sdk/rest_client.py (2)
  • get (197-199)
  • AuthenticationError (46-48)
runagent/sdk/sdk.py (3)
  • configure (67-89)
  • get_config_status (116-118)
  • is_configured (112-114)
runagent/utils/config.py (2)
  • get_config_status (271-288)
  • is_configured (260-268)
runagent/sdk/deployment/middleware_sync.py (4)
  • MiddlewareSyncService (14-224)
  • is_sync_enabled (52-54)
  • _test_middleware_connection (168-177)
  • _test_supabase_authentication (179-188)
runagent/sdk/deployment/middleware_sync.py (2)
runagent/sdk/config.py (2)
  • api_key (280-282)
  • base_url (285-287)
runagent/sdk/rest_client.py (3)
  • RestClient (232-683)
  • RestClient (687-1105)
  • get (197-199)
🔇 Additional comments (3)
templates/agno/default/runagent.config.json (1)

16-19: Extractor for agno_assistant was never configured—no removal occurred
The agno_assistant entrypoint in templates/agno/default/runagent.config.json has always been defined without an extractor. Only the streaming entrypoint (agno_stream) includes an extractor to emit timestamp/content chunks. Since no extractor was ever present on agno_assistant, there are no downstream consumers expecting an extracted payload shape, and no changes are required.

Likely an incorrect or invalid review comment.

runagent/constants.py (1)

19-21: Base URL port bump to 8333 — looks good

Port alignment is consistent with the PR intent and RestClient’s base URL construction. No issues spotted.

Also applies to: 27-27

runagent/sdk/rest_client.py (1)

962-1043: Keep only one validate_api_connection (the UPDATED variant)

This is the better version. After removing the duplicate class/method above, ensure this remains the single implementation inside RestClient.

Run the same grep script from the earlier comment to confirm only one validate_api_connection remains.

Comment on lines +152 to 154
Data,
StructuredData,
}
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Adding MessageType::Data is good; ensure matches and serde round-trips are updated.

Nice parity with Python's MessageType.DATA. Two follow-ups:

  • Verify all match statements over MessageType now handle Data.
  • Add a unit test to assert that MessageType::Data serializes to "data".

Run this to find potential missing match arms:

And consider adding this test:

@@
     fn test_serialization() {
         let entry_point = EntryPoint {
             file: "main.py".to_string(),
             module: "run".to_string(),
             tag: "generic".to_string(),
         };
@@
         assert_eq!(entry_point.tag, deserialized.tag);
     }
+
+    #[test]
+    fn test_message_type_data_serde() {
+        let mt = MessageType::Data;
+        let s = serde_json::to_string(&mt).unwrap();
+        assert_eq!(s, "\"data\"");
+        let back: MessageType = serde_json::from_str(&s).unwrap();
+        assert!(matches!(back, MessageType::Data));
+    }

🏁 Script executed:

#!/bin/bash
# Find matches on MessageType and show context to confirm Data is handled
rg -nP -C3 '\bmatch\s*\(\s*[^)]*MessageType\b|MessageType::(ToolCall|ToolResult|AgentThought|FinalResponse|Error|Status|RawData|StructuredData)\b' --type=rust

Length of output: 5598


Ensure all MessageType::Data arms are implemented and add serde round-trip test

The new MessageType::Data variant currently isn’t handled in your match statements, and there’s no unit test for its serialization. Please address the following:

• Add MessageType::Data arms to every match safe_msg.message_type in
runagent-rust/runagent/src/client/socket_client.rs (around lines 114–124)
– any other match over MessageType in socket_client.rs and
runagent-rust/runagent/src/utils/serializer.rs where you deserialize or handle variants.
• Implement the appropriate logic in those arms (e.g., forwarding Data payloads).
• Add a serde round-trip test in runagent-rust/runagent/src/types/schema.rs:

     #[test]
     fn test_safe_message_with_error() {
         let msg = SafeMessage::with_error("error-id".to_string(), "Something went wrong".to_string());
         assert_eq!(msg.id, "error-id");
         assert!(matches!(msg.message_type, MessageType::Error));
         assert!(msg.error.is_some());
     }

+    #[test]
+    fn test_message_type_data_serde() {
+        let mt = MessageType::Data;
+        let s = serde_json::to_string(&mt).unwrap();
+        assert_eq!(s, "\"data\"");
+        let back: MessageType = serde_json::from_str(&s).unwrap();
+        assert!(matches!(back, MessageType::Data));
+    }

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents
In runagent-rust/runagent/src/types/schema.rs around lines 152–154 and related
modules, the new MessageType::Data variant isn’t handled and lacks a serde
round-trip test; add MessageType::Data match arms in every match over
safe_msg.message_type in runagent-rust/runagent/src/client/socket_client.rs
(around lines 114–124) and any other matches in that file, and update
runagent-rust/runagent/src/utils/serializer.rs to correctly
deserialize/serialize the Data variant (forward the Data payload where
socket_client expects it, preserving payload type and metadata), then add a
serde round-trip unit test in runagent-rust/runagent/src/types/schema.rs that
serializes and deserializes a Message containing MessageType::Data (and one for
StructuredData) and asserts equality to ensure round-trip fidelity.

Comment thread runagent/cli/commands.py
Comment on lines 1611 to +1627

if status or (not enable and not disable and not test):
# Show sync status (default action)
sync_status = sync_service.get_sync_status()

if status or (not test):
# Show detailed sync status
console.print("\n📡 [bold]Middleware Sync Status[/bold]")
console.print("=" * 40)

if sync_status["sync_enabled"]:
console.print("✅ [green]Sync Status: ENABLED[/green]")
# API Key status
if sync_service.api_key:
console.print("🔑 [green]API Key: CONFIGURED[/green]")
console.print(f" Key: [dim]{sync_service.api_key[:16]}...[/dim]")
else:
console.print(" [red]Sync Status: DISABLED[/red]")
console.print("🔑 [red]API Key: NOT CONFIGURED[/red]")

console.print(f"🔑 API Configured: [cyan]{'Yes' if sync_status['api_configured'] else 'No'}[/cyan]")
console.print(f"🌐 Base URL: [blue]{sync_status['base_url']}[/blue]")
# Base URL
console.print(f"🌐 Base URL: [blue]{sync_service.config.base_url}[/blue]")

if sync_status["api_configured"]:
if sync_status["middleware_available"]:
console.print("🟢 [green]Middleware: AVAILABLE[/green]")
else:
console.print("🔴 [red]Middleware: UNAVAILABLE[/red]")
# Authentication status
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Fix inconsistent variable reference auth_validated in local_sync.

The code references sync_service.auth_validated which doesn't exist as a direct attribute. Based on the implementation in middleware_sync.py, this should access the config's internal state.

             # Authentication status
-            if sync_service.auth_validated:
+            if hasattr(sync_service.config, '_config') and sync_service.config._config.get("auth_validated"):
                 console.print("🔐 [green]Authentication: VALID[/green]")
             else:
                 console.print("🔐 [red]Authentication: INVALID[/red]")
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if status or (not enable and not disable and not test):
# Show sync status (default action)
sync_status = sync_service.get_sync_status()
if status or (not test):
# Show detailed sync status
console.print("\n📡 [bold]Middleware Sync Status[/bold]")
console.print("=" * 40)
if sync_status["sync_enabled"]:
console.print("✅ [green]Sync Status: ENABLED[/green]")
# API Key status
if sync_service.api_key:
console.print("🔑 [green]API Key: CONFIGURED[/green]")
console.print(f" Key: [dim]{sync_service.api_key[:16]}...[/dim]")
else:
console.print("❌ [red]Sync Status: DISABLED[/red]")
console.print("🔑 [red]API Key: NOT CONFIGURED[/red]")
console.print(f"🔑 API Configured: [cyan]{'Yes' if sync_status['api_configured'] else 'No'}[/cyan]")
console.print(f"🌐 Base URL: [blue]{sync_status['base_url']}[/blue]")
# Base URL
console.print(f"🌐 Base URL: [blue]{sync_service.config.base_url}[/blue]")
if sync_status["api_configured"]:
if sync_status["middleware_available"]:
console.print("🟢 [green]Middleware: AVAILABLE[/green]")
else:
console.print("🔴 [red]Middleware: UNAVAILABLE[/red]")
# Authentication status
if status or (not test):
# Show detailed sync status
console.print("\n📡 [bold]Middleware Sync Status[/bold]")
console.print("=" * 40)
# API Key status
if sync_service.api_key:
console.print("🔑 [green]API Key: CONFIGURED[/green]")
console.print(f" Key: [dim]{sync_service.api_key[:16]}...[/dim]")
else:
console.print("🔑 [red]API Key: NOT CONFIGURED[/red]")
# Base URL
console.print(f"🌐 Base URL: [blue]{sync_service.config.base_url}[/blue]")
# Authentication status
if hasattr(sync_service.config, '_config') and sync_service.config._config.get("auth_validated"):
console.print("🔐 [green]Authentication: VALID[/green]")
else:
console.print("🔐 [red]Authentication: INVALID[/red]")
🤖 Prompt for AI Agents
In runagent/cli/commands.py around lines 1611 to 1627, the code incorrectly
references sync_service.auth_validated which is not a direct attribute; update
the reference to read the flag from the sync_service.config (e.g.,
sync_service.config.auth_validated or the appropriate accessor on config) so it
checks the config's internal auth validation state, and ensure any downstream
logic uses that config-based value consistently and handles a missing attribute
safely (fallback to False).

Comment on lines +687 to +716
class RestClient:
"""Client for remote server deployment via REST API"""

def __init__(
self,
base_url: Optional[str] = None,
api_key: Optional[str] = None,
api_prefix: Optional[str] = "/api/v1",
):
"""Initialize REST client for middleware server"""
self.api_key = api_key or Config.get_api_key()

# Fix base URL construction
if base_url:
self.base_url = base_url.rstrip("/") + api_prefix
else:
raw_base_url = Config.get_base_url()
self.base_url = raw_base_url.rstrip("/") + api_prefix

# Initialize HTTP handler directly with API key
# The middleware auth system will handle JWT conversion automatically
self.http = HttpHandler(
api_key=self.api_key, # Use API key directly - middleware handles conversion
base_url=self.base_url
)

# Cache for limits to avoid repeated API calls
self._limits_cache = None
self._cache_expiry = None

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Duplicate RestClient class definition — later class overrides the earlier one

A second “class RestClient” redefinition (and a duplicate validate_api_connection) makes the final class omit methods defined in the first class (e.g., get_local_db_limits, upload_agent, deploy_agent, etc.). This will break callers that rely on those methods.

Remove the duplicate class header and the “SIMPLIFIED” method; keep a single RestClient and place only the “UPDATED” validate_api_connection inside it.

Minimal safe patch to remove duplication:

-# runagent/sdk/rest_client.py - FIXED RestClient initialization
-
-class RestClient:
-    """Client for remote server deployment via REST API"""
-
-    def __init__(
-        self,
-        base_url: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_prefix: Optional[str] = "/api/v1",
-    ):
-        """Initialize REST client for middleware server"""
-        self.api_key = api_key or Config.get_api_key()
-        
-        # Fix base URL construction
-        if base_url:
-            self.base_url = base_url.rstrip("/") + api_prefix
-        else:
-            raw_base_url = Config.get_base_url()
-            self.base_url = raw_base_url.rstrip("/") + api_prefix
-
-        # Initialize HTTP handler directly with API key
-        # The middleware auth system will handle JWT conversion automatically
-        self.http = HttpHandler(
-            api_key=self.api_key,  # Use API key directly - middleware handles conversion
-            base_url=self.base_url
-        )
-
-        # Cache for limits to avoid repeated API calls
-        self._limits_cache = None
-        self._cache_expiry = None
-
-    def validate_api_connection(self) -> Dict[str, Any]:
-        """Validate API connection and authentication - SIMPLIFIED"""
-        try:
-            # Test basic connectivity first
-            try:
-                health_response = self.http.get("/health", timeout=10, handle_errors=False)
-                
-                if health_response.status_code != 200:
-                    return {
-                        "success": False,
-                        "api_connected": False,
-                        "error": f"Health check failed: {health_response.status_code}",
-                    }
-
-            except Exception as e:
-                return {
-                    "success": False,
-                    "api_connected": False,
-                    "error": f"Cannot connect to middleware: {str(e)}",
-                }
-
-            # Test authentication if API key provided
-            if self.api_key:
-                try:
-                    # Try to get user profile which requires authentication
-                    auth_response = self.http.get("/users/profile", timeout=10)
-                    
-                    if auth_response.status_code == 200:
-                        profile_data = auth_response.json()
-                        auth_data = profile_data.get("auth_data", {})
-                        
-                        return {
-                            "success": True,
-                            "api_connected": True,
-                            "api_authenticated": True,
-                            "user_info": {
-                                "email": auth_data.get("email"),
-                                "id": auth_data.get("id")
-                            },
-                            "base_url": self.base_url,
-                        }
-                    else:
-                        error_data = auth_response.json() if hasattr(auth_response, 'json') else {}
-                        return {
-                            "success": False,
-                            "api_connected": True,
-                            "api_authenticated": False,
-                            "error": error_data.get("detail", f"Authentication failed: {auth_response.status_code}"),
-                            "base_url": self.base_url,
-                        }
-                        
-                except AuthenticationError as e:
-                    return {
-                        "success": False,
-                        "api_connected": True,
-                        "api_authenticated": False,
-                        "error": f"Invalid API key: {e.message}",
-                        "base_url": self.base_url,
-                    }
-                except Exception as e:
-                    return {
-                        "success": False,
-                        "api_connected": True,
-                        "api_authenticated": False,
-                        "error": f"Authentication test failed: {str(e)}",
-                        "base_url": self.base_url,
-                    }
-            else:
-                return {
-                    "success": True,
-                    "api_connected": True,
-                    "api_authenticated": False,
-                    "base_url": self.base_url,
-                    "message": "No API key provided",
-                }
-
-        except Exception as e:
-            return {
-                "success": False,
-                "api_connected": False,
-                "error": f"Connection validation failed: {str(e)}",
-            }

Then, add the “UPDATED” validate_api_connection once inside the original class (placement right after init is fine). Suggested insertion (outside the changed hunk, showing only for clarity):

def validate_api_connection(self) -> Dict[str, Any]:
    """Validate API connection and authentication - UPDATED"""
    # ... use the implementation from lines 962–1043 (unchanged) ...

Also applies to: 718-792


🏁 Script executed:

#!/bin/bash
# Show duplicates to ensure only one class remains after your edits
rg -nP '^\s*class\s+RestClient\b' runagent/sdk/rest_client.py -n
# Grep both validate_api_connection definitions
rg -nP '^\s*def\s+validate_api_connection\s*\(' runagent/sdk/rest_client.py -n -C2

Length of output: 542


Remove duplicate RestClient and consolidate validate_api_connection

The file defines RestClient twice—with a “SIMPLIFIED” validate_api_connection in the second class (lines 687–961) overriding all methods from the first (line 232). Consolidate into the original class by removing the duplicate definition and merging only the “UPDATED” validation logic.

• Duplicate class declarations at:
– First RestClient starts at line 232 (includes methods like get_local_db_limits, upload_agent, deploy_agent, etc.)
– Second RestClient starts at line 687 with a simplified validate_api_connection (lines 718–961)
• Two validate_api_connection methods:
– SIMPLIFIED version at lines 718–961 (inside duplicate class)
– UPDATED version at lines 962–1043
• Action:

  1. Delete lines 687–961 (the entire duplicate class and its SIMPLIFIED validation).
  2. Insert the UPDATED validate_api_connection (from lines 962–1043) into the first RestClient (e.g. immediately after its init).

Minimal patch sketch:

--- a/runagent/sdk/rest_client.py
+++ b/runagent/sdk/rest_client.py
@@ -687,275 +687,0 @@
-class RestClient:
-    """Client for remote server deployment via REST API"""
-    def __init__(...):
-        ...
-    def validate_api_connection(self) -> Dict[str, Any]:
-        """Validate API connection and authentication – SIMPLIFIED"""
-        ...
-        # entire duplicate class removed
@@ -232,6 +232,47 @@ class RestClient:
     def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None, api_prefix: Optional[str] = "/api/v1"):
         ...
+
+    def validate_api_connection(self) -> Dict[str, Any]:
+        """Validate API connection and authentication – UPDATED"""
+        try:
+            # Test basic connectivity
+            health_response = self.http.get("/health", timeout=10, handle_errors=False)
+            if health_response.status_code != 200:
+                return {
+                    "success": False,
+                    "api_connected": False,
+                    "error": f"Health check failed: {health_response.status_code}",
+                }
+        except Exception as e:
+            return {
+                "success": False,
+                "api_connected": False,
+                "error": f"Cannot connect to middleware: {str(e)}",
+            }
+
+        if self.api_key:
+            # Authentication flow (unchanged from lines 962–1043)
+            try:
+                auth_response = self.http.get("/users/profile", timeout=10)
+                # … rest of UPDATED logic …
+            except AuthenticationError as e:
+                # … handle invalid key …
+            except Exception as e:
+                # … handle other errors …
+        else:
+            return {
+                "success": True,
+                "api_connected": True,
+                "api_authenticated": False,
+                "base_url": self.base_url,
+                "message": "No API key provided",
+            }

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents
In runagent/sdk/rest_client.py around lines 232 and 687–1043: the file contains
a duplicate RestClient class (second starts at line 687) with a SIMPLIFIED
validate_api_connection that overrides the original; remove the entire duplicate
class block (delete lines 687–961) and copy the UPDATED validate_api_connection
implementation (lines 962–1043) into the original RestClient class (place it
immediately after the __init__ method), ensuring there is only one RestClient
definition with the updated validation logic and no leftover duplicate methods
or class declarations.

Comment on lines +125 to 127
"server_host": "127.0.0.1", # Add proper server info
"server_port": 8450 # Add proper port info
}
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Extract hardcoded server configuration values.

The server host and port are hardcoded. These should be passed as parameters or retrieved from configuration.

                         "client_info": {
                             "connection_type": "websocket",
                             "stream_mode": True,
-                            "server_host": "127.0.0.1",  # Add proper server info
-                            "server_port": 8450  # Add proper port info
+                            "server_host": getattr(self, 'server_host', "127.0.0.1"),
+                            "server_port": getattr(self, 'server_port', 8450)
                         }

Consider adding these as initialization parameters:

-    def __init__(self, db_service, middleware_sync=None):
+    def __init__(self, db_service, middleware_sync=None, server_host="127.0.0.1", server_port=8450):
         self.db_service = db_service
         self.serializer = CoreSerializer(max_size_mb=10.0)
         self.middleware_sync = middleware_sync or get_middleware_sync()
         self.active_streams = {}
+        self.server_host = server_host
+        self.server_port = server_port
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
"server_host": "127.0.0.1", # Add proper server info
"server_port": 8450 # Add proper port info
}
"client_info": {
"connection_type": "websocket",
"stream_mode": True,
"server_host": getattr(self, 'server_host', "127.0.0.1"),
"server_port": getattr(self, 'server_port', 8450)
}
Suggested change
"server_host": "127.0.0.1", # Add proper server info
"server_port": 8450 # Add proper port info
}
def __init__(self, db_service, middleware_sync=None, server_host="127.0.0.1", server_port=8450):
self.db_service = db_service
self.serializer = CoreSerializer(max_size_mb=10.0)
self.middleware_sync = middleware_sync or get_middleware_sync()
self.active_streams = {}
self.server_host = server_host
self.server_port = server_port
🤖 Prompt for AI Agents
In runagent/sdk/server/socket_utils.py around lines 125 to 127, the server_host
and server_port are hardcoded ("127.0.0.1" and 8450); change them to be
configurable by extracting them into initialization parameters (e.g.,
SocketUtils.__init__(server_host, server_port, ...) or factory function) or read
from a central config/env with sensible defaults, validate/normalize the values
(host string and integer port range) and replace the literal dict entries with
the injected/config values; update all callers/tests to pass the new parameters
or ensure config is available.

@sawradip sawradip merged commit b55f297 into main Aug 23, 2025
2 checks passed
This was referenced Oct 7, 2025
@coderabbitai coderabbitai Bot mentioned this pull request Nov 12, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants