From 27b1066e1b4beb4a93c12813e745760033e92d6a Mon Sep 17 00:00:00 2001 From: Jeff Zeng Date: Tue, 10 Feb 2026 18:28:55 +1100 Subject: [PATCH] Fix critical bugs preventing CLI execution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Fix --query CLI argument not working - The argument was parsed but never used, always running interactive mode - Added conditional to run single-query mode when args.query is provided 2. Fix session not waiting for READY state - After start_code_interpreter_session(), code tried to use session immediately while still in "STARTING" state - Added polling loop (up to 2 min) to wait for READY status 3. Fix Python module namespace bug in utils.py - When running as script (__main__), globals weren't visible to utils.py because it imported by module name (creating separate module instance) - Changed to sys.modules['__main__'] to access actual running module globals 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- code_exec_with_mcp_agent.py | 79 ++++++++++++++++++++++++++++++++++++- utils.py | 15 ++++--- 2 files changed, 86 insertions(+), 8 deletions(-) diff --git a/code_exec_with_mcp_agent.py b/code_exec_with_mcp_agent.py index 42572ca..13d9891 100644 --- a/code_exec_with_mcp_agent.py +++ b/code_exec_with_mcp_agent.py @@ -191,6 +191,31 @@ def _get_or_create_code_client( _session_id = session_response["sessionId"] logger.info(f"Started new session: {_session_id}") + # Wait for session to become READY + import time + max_wait = 120 # 2 minutes max + poll_interval = 2 # seconds + waited = 0 + while waited < max_wait: + try: + resp = _dp_client.get_code_interpreter_session( + codeInterpreterIdentifier=_interpreter_id, + sessionId=_session_id, + ) + status = resp.get("status") + logger.info(f"Session status: {status} (waited {waited}s)") + if status == "READY": + logger.info("Session is READY") + break + elif status in ["FAILED", "STOPPED", "STOPPING"]: + raise RuntimeError(f"Session failed to start: {status}") + except Exception as e: + logger.warning(f"Error checking session status: {e}") + time.sleep(poll_interval) + waited += poll_interval + else: + raise RuntimeError(f"Session did not become READY within {max_wait}s") + # Reset per-session flags for new session _env_initialized = False _files_uploaded = False @@ -238,6 +263,31 @@ def _get_or_create_code_client( _session_id = session_response["sessionId"] logger.info(f"Started new session: {_session_id}") + # Wait for session to become READY + import time + max_wait = 120 # 2 minutes max + poll_interval = 2 # seconds + waited = 0 + while waited < max_wait: + try: + resp = _dp_client.get_code_interpreter_session( + codeInterpreterIdentifier=_interpreter_id, + sessionId=_session_id, + ) + status = resp.get("status") + logger.info(f"Session status: {status} (waited {waited}s)") + if status == "READY": + logger.info("Session is READY") + break + elif status in ["FAILED", "STOPPED", "STOPPING"]: + raise RuntimeError(f"Session failed to start: {status}") + except Exception as e: + logger.warning(f"Error checking session status: {e}") + time.sleep(poll_interval) + waited += poll_interval + else: + raise RuntimeError(f"Session did not become READY within {max_wait}s") + # New sandbox → reset per-session flags _env_initialized = False _files_uploaded = False @@ -1100,8 +1150,33 @@ def main() -> None: print(f"\nError creating agent: {e}") sys.exit(1) - # Run agent in interactive mode - _run_interactive_mode(agent, checkpointer, model_id=args.model) + # Run agent in single query or interactive mode + if args.query: + # Single query mode + logger.info(f"Running single query: {args.query}") + thread_id = "single_query_thread" + config = {"configurable": {"thread_id": thread_id}} + + print("\n" + "-" * 80) + print("Response:") + print("-" * 80) + + result = _execute_agent_query(agent, args.query, config, args.model) + + # Print response + print(result["response"]) + print() + + # Print metrics + if result["total_tokens"] > 0: + print(f"\nMetrics:") + print(f" Latency: {result['latency_seconds']:.2f}s") + print(f" Tokens: {result['total_tokens']:,} " + f"(in: {result['input_tokens']:,}, out: {result['output_tokens']:,})") + print(f" Cost: ${result['estimated_cost_usd']:.6f}") + else: + # Interactive mode + _run_interactive_mode(agent, checkpointer, model_id=args.model) if __name__ == "__main__": diff --git a/utils.py b/utils.py index e2aa1e1..70a302b 100644 --- a/utils.py +++ b/utils.py @@ -140,12 +140,15 @@ def call_tool( Example: result = call_tool("executeCommand", {"command": "pip install requests"}) """ - # Import globals from code_exec_with_mcp_agent module - import code_exec_with_mcp_agent - - _dp_client = code_exec_with_mcp_agent._dp_client - _interpreter_id = code_exec_with_mcp_agent._interpreter_id - _session_id = code_exec_with_mcp_agent._session_id + # Import globals from the main module + # Note: When running as a script, the module is __main__, not code_exec_with_mcp_agent + # So we need to use sys.modules to get the actual running module + import sys + main_module = sys.modules.get('__main__') + + _dp_client = getattr(main_module, '_dp_client', None) + _interpreter_id = getattr(main_module, '_interpreter_id', None) + _session_id = getattr(main_module, '_session_id', None) if not _dp_client or not _interpreter_id or not _session_id: raise RuntimeError(