netra-zen 1.1.2__tar.gz → 1.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (76) hide show
  1. {netra_zen-1.1.2/netra_zen.egg-info → netra_zen-1.2.0}/PKG-INFO +1 -1
  2. {netra_zen-1.1.2 → netra_zen-1.2.0/netra_zen.egg-info}/PKG-INFO +1 -1
  3. {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/SOURCES.txt +1 -0
  4. {netra_zen-1.1.2 → netra_zen-1.2.0}/pyproject.toml +1 -1
  5. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/agent_cli.py +108 -9
  6. {netra_zen-1.1.2 → netra_zen-1.2.0}/setup.py +1 -1
  7. netra_zen-1.2.0/tests/test_timing_fix.py +243 -0
  8. {netra_zen-1.1.2 → netra_zen-1.2.0}/LICENSE.md +0 -0
  9. {netra_zen-1.1.2 → netra_zen-1.2.0}/MANIFEST.in +0 -0
  10. {netra_zen-1.1.2 → netra_zen-1.2.0}/README.md +0 -0
  11. {netra_zen-1.1.2 → netra_zen-1.2.0}/agent_interface/__init__.py +0 -0
  12. {netra_zen-1.1.2 → netra_zen-1.2.0}/agent_interface/base_agent.py +0 -0
  13. {netra_zen-1.1.2 → netra_zen-1.2.0}/config_example.json +0 -0
  14. {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/CACHE_TOKENS_GUIDE.md +0 -0
  15. {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/Cost_allocation.md +0 -0
  16. {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/DOLLAR_BUDGET_USAGE_EXAMPLES.md +0 -0
  17. {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/EXAMPLES.md +0 -0
  18. {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/MODEL_COLUMN_GUIDE.md +0 -0
  19. {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/apex_integration_test_plan.md +0 -0
  20. {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/zen_agent_cli_parallel_plan.md +0 -0
  21. {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/dependency_links.txt +0 -0
  22. {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/entry_points.txt +0 -0
  23. {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/requires.txt +0 -0
  24. {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/top_level.txt +0 -0
  25. {netra_zen-1.1.2 → netra_zen-1.2.0}/prebuilt_commands_example.json +0 -0
  26. {netra_zen-1.1.2 → netra_zen-1.2.0}/requirements-dev.txt +0 -0
  27. {netra_zen-1.1.2 → netra_zen-1.2.0}/requirements.txt +0 -0
  28. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/__init__.py +0 -0
  29. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/__main__.py +0 -0
  30. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/agent_logs.py +0 -0
  31. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/bump_version.py +0 -0
  32. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/demo_log_collection.py +0 -0
  33. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/embed_release_credentials.py +0 -0
  34. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/test_apex_telemetry_debug.py +0 -0
  35. {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/verify_log_transmission.py +0 -0
  36. {netra_zen-1.1.2 → netra_zen-1.2.0}/setup.cfg +0 -0
  37. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/__init__.py +0 -0
  38. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/htmlcov/status.json +0 -0
  39. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_agent_interface.py +0 -0
  40. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_agent_logs.py +0 -0
  41. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_integration.py +0 -0
  42. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_streaming_fix.py +0 -0
  43. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry.py +0 -0
  44. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry_live.py +0 -0
  45. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry_mock.py +0 -0
  46. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry_regression.py +0 -0
  47. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_cli_extensions.py +0 -0
  48. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_cli_integration.py +0 -0
  49. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_direct_command_execution.py +0 -0
  50. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_dollar_budget_enhancement.py +0 -0
  51. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_event_batching_issue.py +0 -0
  52. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_permission_fix_windows.py +0 -0
  53. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_pricing_engine.py +0 -0
  54. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_runner.py +0 -0
  55. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_handshake.py +0 -0
  56. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_id_fix.py +0 -0
  57. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_management.py +0 -0
  58. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_resolution.py +0 -0
  59. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_workspace_detection.py +0 -0
  60. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_commands.py +0 -0
  61. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_integration.py +0 -0
  62. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_metrics.py +0 -0
  63. {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_unit.py +0 -0
  64. {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/__init__.py +0 -0
  65. {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/budget_manager.py +0 -0
  66. {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/models.py +0 -0
  67. {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/visualization.py +0 -0
  68. {netra_zen-1.1.2 → netra_zen-1.2.0}/token_transparency/__init__.py +0 -0
  69. {netra_zen-1.1.2 → netra_zen-1.2.0}/token_transparency/claude_pricing_engine.py +0 -0
  70. {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/__init__.py +0 -0
  71. {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/__main__.py +0 -0
  72. {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/__init__.py +0 -0
  73. {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/apex_telemetry.py +0 -0
  74. {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/embedded_credentials.py +0 -0
  75. {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/manager.py +0 -0
  76. {netra_zen-1.1.2 → netra_zen-1.2.0}/zen_orchestrator.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: netra-zen
3
- Version: 1.1.2
3
+ Version: 1.2.0
4
4
  Summary: Multi-instance Claude orchestrator for parallel task execution
5
5
  Home-page: https://github.com/netra-systems/zen
6
6
  Author: Systems
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: netra-zen
3
- Version: 1.1.2
3
+ Version: 1.2.0
4
4
  Summary: Multi-instance Claude orchestrator for parallel task execution
5
5
  Home-page: https://github.com/netra-systems/zen
6
6
  Author: Systems
@@ -53,6 +53,7 @@ tests/test_thread_handshake.py
53
53
  tests/test_thread_id_fix.py
54
54
  tests/test_thread_management.py
55
55
  tests/test_thread_resolution.py
56
+ tests/test_timing_fix.py
56
57
  tests/test_workspace_detection.py
57
58
  tests/test_zen_commands.py
58
59
  tests/test_zen_integration.py
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "netra-zen"
7
- version = "1.1.2"
7
+ version = "1.2.0"
8
8
  description = "Multi-instance Claude orchestrator for parallel task execution"
9
9
  readme = "README.md"
10
10
  authors = [
@@ -2883,22 +2883,49 @@ class WebSocketClient:
2883
2883
  self.connected = True
2884
2884
  return True
2885
2885
  else:
2886
- # Handshake failed - backend might be using old reactive model
2886
+ # Handshake failed - server not ready to process messages
2887
2887
  self.debug.debug_print(
2888
- "WARNING: No proactive handshake received from server",
2888
+ "WARNING: Server handshake not completed - server not ready",
2889
2889
  DebugLevel.BASIC,
2890
2890
  style="yellow"
2891
2891
  )
2892
2892
  self.debug.debug_print(
2893
- "Server may be using pre-2025-10-09 architecture (reactive handshake)",
2893
+ "Server still completing lifecycle phases (Initialize → Authenticate → Handshake → Prepare → Processing)",
2894
2894
  DebugLevel.VERBOSE,
2895
2895
  style="yellow"
2896
2896
  )
2897
2897
 
2898
- # Still mark as connected for backward compatibility
2899
- # Old backends might work without the proactive handshake
2900
- self.connected = True
2901
- return True
2898
+ # Wait briefly for server to complete its phases and retry
2899
+ safe_console_print("⏳ Waiting for server to complete initialization...", style="yellow",
2900
+ json_mode=self.config.json_mode, ci_mode=self.config.ci_mode)
2901
+ await asyncio.sleep(3.0) # Give server time to reach PROCESSING phase
2902
+
2903
+ # Retry handshake once more
2904
+ self.debug.debug_print(
2905
+ "Retrying handshake after delay...",
2906
+ DebugLevel.VERBOSE,
2907
+ style="cyan"
2908
+ )
2909
+ handshake_success = await self._perform_handshake()
2910
+ if handshake_success:
2911
+ safe_console_print(f"✅ Connected with thread ID: {self.current_thread_id}", style="green",
2912
+ json_mode=self.config.json_mode, ci_mode=self.config.ci_mode)
2913
+ self.connected = True
2914
+ return True
2915
+ else:
2916
+ # Server still not ready - fail the connection
2917
+ self.debug.debug_print(
2918
+ "ERROR: Server not ready after retry",
2919
+ DebugLevel.BASIC,
2920
+ style="red"
2921
+ )
2922
+ safe_console_print("❌ Server not ready to process messages. Please try again.", style="red",
2923
+ json_mode=self.config.json_mode, ci_mode=self.config.ci_mode)
2924
+
2925
+ # Close WebSocket and fail gracefully
2926
+ if self.ws:
2927
+ await self.ws.close()
2928
+ return False
2902
2929
  except Exception as e:
2903
2930
  self.debug.log_connection_attempt(method_name, self.config.ws_url, success=False, error=str(e))
2904
2931
  self.debug.debug_print(
@@ -3027,7 +3054,8 @@ class WebSocketClient:
3027
3054
 
3028
3055
  # Server enters HANDSHAKING phase after authentication
3029
3056
  # and proactively sends handshake_response
3030
- handshake_timeout = 10.0 # Wait up to 10 seconds
3057
+ # Use configured timeout or default to 10 seconds
3058
+ handshake_timeout = self.handshake_timeout if hasattr(self, 'handshake_timeout') and self.handshake_timeout else 10.0
3031
3059
  start_time = asyncio.get_event_loop().time()
3032
3060
 
3033
3061
  try:
@@ -3274,14 +3302,64 @@ class WebSocketClient:
3274
3302
  )
3275
3303
 
3276
3304
  # CRITICAL: Send acknowledgment with the SAME thread_id
3305
+ # Per WebSocket Client Lifecycle Guide, use "handshake_acknowledged"
3277
3306
  ack_message = {
3278
- "type": "session_acknowledged",
3307
+ "type": "handshake_acknowledged",
3279
3308
  "thread_id": backend_thread_id, # Echo back the same ID
3280
3309
  "timestamp": datetime.now(timezone.utc).isoformat()
3281
3310
  }
3282
3311
 
3283
3312
  await self.ws.send(json.dumps(ack_message))
3284
3313
 
3314
+ # Per WebSocket Client Lifecycle Guide, wait for handshake_complete
3315
+ # and then add delay for server to enter Phase 5 (Processing)
3316
+ self.debug.debug_print(
3317
+ "Waiting for handshake_complete confirmation...",
3318
+ DebugLevel.VERBOSE,
3319
+ style="cyan"
3320
+ )
3321
+
3322
+ # Wait briefly for handshake_complete message
3323
+ try:
3324
+ complete_msg = await asyncio.wait_for(self.ws.recv(), timeout=2.0)
3325
+ complete_data = json.loads(complete_msg)
3326
+ if complete_data.get('type') == 'handshake_complete':
3327
+ self.debug.debug_print(
3328
+ "✅ Received handshake_complete - Server at Phase 4 (Ready)",
3329
+ DebugLevel.VERBOSE,
3330
+ style="green"
3331
+ )
3332
+
3333
+ # CRITICAL: Add delay for server to enter Phase 5 (Processing)
3334
+ # Per documentation, 500ms is recommended
3335
+ self.debug.debug_print(
3336
+ "Waiting 500ms for server to enter Phase 5 (Processing)...",
3337
+ DebugLevel.VERBOSE,
3338
+ style="cyan"
3339
+ )
3340
+ await asyncio.sleep(0.5)
3341
+
3342
+ self.debug.debug_print(
3343
+ "✅ Server should now be in Phase 5 (Processing) - ready for messages",
3344
+ DebugLevel.VERBOSE,
3345
+ style="green"
3346
+ )
3347
+ except asyncio.TimeoutError:
3348
+ # If no handshake_complete, still add a delay to be safe
3349
+ self.debug.debug_print(
3350
+ "No handshake_complete received, adding safety delay",
3351
+ DebugLevel.VERBOSE,
3352
+ style="yellow"
3353
+ )
3354
+ await asyncio.sleep(0.5)
3355
+ except Exception as e:
3356
+ self.debug.debug_print(
3357
+ f"Error waiting for handshake_complete: {e}",
3358
+ DebugLevel.VERBOSE,
3359
+ style="yellow"
3360
+ )
3361
+ await asyncio.sleep(0.5)
3362
+
3285
3363
  return True
3286
3364
 
3287
3365
  def _get_platform_cache_path(self) -> Path:
@@ -3622,6 +3700,27 @@ class WebSocketClient:
3622
3700
  if not self.ws:
3623
3701
  raise RuntimeError("WebSocket not connected")
3624
3702
 
3703
+ # Check if connection handshake is fully complete
3704
+ if not self.connected:
3705
+ self.debug.debug_print(
3706
+ "ERROR: Connection not ready - handshake not complete",
3707
+ DebugLevel.BASIC,
3708
+ style="red"
3709
+ )
3710
+ safe_console_print(
3711
+ "\n❌ ERROR: Cannot send message - server not ready",
3712
+ style="red",
3713
+ json_mode=self.config.json_mode,
3714
+ ci_mode=self.config.ci_mode
3715
+ )
3716
+ safe_console_print(
3717
+ "The server is still completing its initialization phases.",
3718
+ style="yellow",
3719
+ json_mode=self.config.json_mode,
3720
+ ci_mode=self.config.ci_mode
3721
+ )
3722
+ raise RuntimeError("Connection not ready - handshake incomplete. Server needs to reach Processing phase.")
3723
+
3625
3724
  # Generate run_id
3626
3725
  self.run_id = f"cli_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{os.getpid()}"
3627
3726
 
@@ -8,7 +8,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
8
8
 
9
9
  setup(
10
10
  name="netra-zen",
11
- version="1.1.2",
11
+ version="1.2.0",
12
12
  author=" Systems",
13
13
  author_email="pypi@netrasystems.ai",
14
14
  description="Multi-instance Claude orchestrator for parallel task execution",
@@ -0,0 +1,243 @@
1
+ """
2
+ Test for the timing fix that ensures CLI waits for server handshake
3
+ before sending messages.
4
+
5
+ This test verifies that:
6
+ 1. CLI waits for handshake_response before marking connection as ready
7
+ 2. CLI retries handshake with delay if server isn't ready
8
+ 3. CLI properly fails connection if server doesn't complete handshake
9
+ 4. No messages are sent before server reaches PROCESSING phase
10
+ """
11
+
12
+ import asyncio
13
+ import json
14
+ import unittest
15
+ from unittest.mock import AsyncMock, MagicMock, patch, call
16
+ from datetime import datetime
17
+ import sys
18
+ import os
19
+
20
+ # Add parent directory to path for imports
21
+ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
22
+
23
+ from scripts.agent_cli import WebSocketClient
24
+ from shared import Environment, EnvironmentConfig
25
+
26
+
27
+ class TestTimingFix(unittest.TestCase):
28
+ """Test suite for CLI timing fix"""
29
+
30
+ def setUp(self):
31
+ """Set up test fixtures"""
32
+ self.config = MagicMock()
33
+ self.config.ws_url = "ws://localhost:8000/ws"
34
+ self.config.environment = Environment.LOCAL
35
+ self.config.json_mode = False
36
+ self.config.ci_mode = False
37
+ self.token = "test_token"
38
+ self.debug = MagicMock()
39
+
40
+ @patch('scripts.agent_cli.websockets.connect')
41
+ async def test_successful_handshake_on_first_try(self, mock_connect):
42
+ """Test successful handshake on first attempt"""
43
+ # Mock WebSocket connection
44
+ mock_ws = AsyncMock()
45
+ mock_connect.return_value = mock_ws
46
+
47
+ # Mock handshake response
48
+ handshake_response = {
49
+ "type": "handshake_response",
50
+ "thread_id": "test-thread-123",
51
+ "user_id": "user-456",
52
+ "session_id": "session-789"
53
+ }
54
+ mock_ws.recv.return_value = json.dumps(handshake_response)
55
+
56
+ # Create client and connect
57
+ client = WebSocketClient(
58
+ self.config, self.token, self.debug,
59
+ send_logs=False, logs_count=10,
60
+ logs_project=None, logs_path=None,
61
+ logs_user=None, handshake_timeout=5
62
+ )
63
+
64
+ result = await client.connect()
65
+
66
+ # Assertions
67
+ self.assertTrue(result)
68
+ self.assertTrue(client.connected)
69
+ self.assertEqual(client.current_thread_id, "test-thread-123")
70
+ mock_ws.close.assert_not_called()
71
+
72
+ @patch('scripts.agent_cli.websockets.connect')
73
+ @patch('scripts.agent_cli.asyncio.sleep')
74
+ async def test_handshake_retry_after_delay(self, mock_sleep, mock_connect):
75
+ """Test handshake retry after delay when server not ready"""
76
+ # Mock WebSocket connection
77
+ mock_ws = AsyncMock()
78
+ mock_connect.return_value = mock_ws
79
+
80
+ # First attempt: timeout (server not ready)
81
+ # Second attempt: successful handshake
82
+ handshake_response = {
83
+ "type": "handshake_response",
84
+ "thread_id": "test-thread-123",
85
+ "user_id": "user-456",
86
+ "session_id": "session-789"
87
+ }
88
+
89
+ # Configure recv to timeout first, then succeed
90
+ mock_ws.recv.side_effect = [
91
+ asyncio.TimeoutError(), # First handshake attempt times out
92
+ json.dumps(handshake_response) # Second attempt succeeds
93
+ ]
94
+
95
+ # Create client and connect
96
+ client = WebSocketClient(
97
+ self.config, self.token, self.debug,
98
+ send_logs=False, logs_count=10,
99
+ logs_project=None, logs_path=None,
100
+ logs_user=None, handshake_timeout=1
101
+ )
102
+
103
+ result = await client.connect()
104
+
105
+ # Assertions
106
+ self.assertTrue(result)
107
+ self.assertTrue(client.connected)
108
+ self.assertEqual(client.current_thread_id, "test-thread-123")
109
+ mock_sleep.assert_called_once_with(3.0) # Verify delay was applied
110
+ mock_ws.close.assert_not_called()
111
+
112
+ @patch('scripts.agent_cli.websockets.connect')
113
+ @patch('scripts.agent_cli.asyncio.sleep')
114
+ async def test_connection_fails_if_no_handshake(self, mock_sleep, mock_connect):
115
+ """Test connection fails properly when server never completes handshake"""
116
+ # Mock WebSocket connection
117
+ mock_ws = AsyncMock()
118
+ mock_connect.return_value = mock_ws
119
+
120
+ # Both handshake attempts timeout
121
+ mock_ws.recv.side_effect = [
122
+ asyncio.TimeoutError(), # First attempt
123
+ asyncio.TimeoutError() # Second attempt after retry
124
+ ]
125
+
126
+ # Create client and connect
127
+ client = WebSocketClient(
128
+ self.config, self.token, self.debug,
129
+ send_logs=False, logs_count=10,
130
+ logs_project=None, logs_path=None,
131
+ logs_user=None, handshake_timeout=1
132
+ )
133
+
134
+ result = await client.connect()
135
+
136
+ # Assertions
137
+ self.assertFalse(result) # Connection should fail
138
+ self.assertFalse(client.connected) # Should not be marked as connected
139
+ self.assertIsNone(client.current_thread_id) # No thread ID assigned
140
+ mock_sleep.assert_called_once_with(3.0) # Verify retry delay
141
+ mock_ws.close.assert_called_once() # WebSocket should be closed
142
+
143
+ @patch('scripts.agent_cli.websockets.connect')
144
+ async def test_no_message_sent_without_handshake(self, mock_connect):
145
+ """Test that send_message fails if handshake not completed"""
146
+ # Mock WebSocket connection
147
+ mock_ws = AsyncMock()
148
+ mock_connect.return_value = mock_ws
149
+
150
+ # Handshake times out - no response
151
+ mock_ws.recv.side_effect = asyncio.TimeoutError()
152
+
153
+ # Create client and try to connect
154
+ client = WebSocketClient(
155
+ self.config, self.token, self.debug,
156
+ send_logs=False, logs_count=10,
157
+ logs_project=None, logs_path=None,
158
+ logs_user=None, handshake_timeout=1
159
+ )
160
+
161
+ # Connection should fail due to no handshake
162
+ result = await client.connect()
163
+ self.assertFalse(result)
164
+
165
+ # Attempting to send message should fail
166
+ with self.assertRaises(Exception) as context:
167
+ await client.send_message("test message")
168
+
169
+ # Verify no message was sent via WebSocket
170
+ mock_ws.send.assert_not_called()
171
+
172
+ @patch('scripts.agent_cli.websockets.connect')
173
+ async def test_connection_established_is_not_handshake(self, mock_connect):
174
+ """Test that connection_established event is not treated as handshake"""
175
+ # Mock WebSocket connection
176
+ mock_ws = AsyncMock()
177
+ mock_connect.return_value = mock_ws
178
+
179
+ # Send connection_established instead of handshake_response
180
+ connection_event = {
181
+ "type": "connection_established",
182
+ "message": "Connected to server"
183
+ }
184
+
185
+ # Then timeout (no actual handshake)
186
+ mock_ws.recv.side_effect = [
187
+ json.dumps(connection_event),
188
+ asyncio.TimeoutError()
189
+ ]
190
+
191
+ # Create client and connect
192
+ client = WebSocketClient(
193
+ self.config, self.token, self.debug,
194
+ send_logs=False, logs_count=10,
195
+ logs_project=None, logs_path=None,
196
+ logs_user=None, handshake_timeout=1
197
+ )
198
+
199
+ result = await client.connect()
200
+
201
+ # Connection should fail - connection_established is not a handshake
202
+ self.assertFalse(result)
203
+ self.assertFalse(client.connected)
204
+ self.assertIsNone(client.current_thread_id)
205
+
206
+
207
+ def run_async_test(coro):
208
+ """Helper to run async tests"""
209
+ loop = asyncio.get_event_loop()
210
+ return loop.run_until_complete(coro)
211
+
212
+
213
+ class AsyncTestRunner(unittest.TestCase):
214
+ """Wrapper to run async tests"""
215
+
216
+ def test_successful_handshake(self):
217
+ test = TestTimingFix()
218
+ test.setUp()
219
+ run_async_test(test.test_successful_handshake_on_first_try())
220
+
221
+ def test_retry_after_delay(self):
222
+ test = TestTimingFix()
223
+ test.setUp()
224
+ run_async_test(test.test_handshake_retry_after_delay())
225
+
226
+ def test_connection_fails(self):
227
+ test = TestTimingFix()
228
+ test.setUp()
229
+ run_async_test(test.test_connection_fails_if_no_handshake())
230
+
231
+ def test_no_early_messages(self):
232
+ test = TestTimingFix()
233
+ test.setUp()
234
+ run_async_test(test.test_no_message_sent_without_handshake())
235
+
236
+ def test_connection_established_ignored(self):
237
+ test = TestTimingFix()
238
+ test.setUp()
239
+ run_async_test(test.test_connection_established_is_not_handshake())
240
+
241
+
242
+ if __name__ == '__main__':
243
+ unittest.main()
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes