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(