waldiez 0.5.8__py3-none-any.whl → 0.5.10__py3-none-any.whl
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.
Potentially problematic release.
This version of waldiez might be problematic. Click here for more details.
- waldiez/_version.py +1 -1
- waldiez/cli.py +112 -24
- waldiez/exporting/agent/exporter.py +3 -0
- waldiez/exporting/agent/extras/captain_agent_extras.py +44 -7
- waldiez/exporting/agent/extras/handoffs/condition.py +3 -1
- waldiez/exporting/chats/utils/common.py +25 -23
- waldiez/exporting/core/__init__.py +0 -2
- waldiez/exporting/core/context.py +13 -13
- waldiez/exporting/core/protocols.py +0 -141
- waldiez/exporting/core/result.py +5 -5
- waldiez/exporting/flow/merger.py +2 -2
- waldiez/exporting/flow/orchestrator.py +1 -0
- waldiez/exporting/flow/utils/common.py +2 -2
- waldiez/exporting/flow/utils/importing.py +1 -0
- waldiez/exporting/flow/utils/logging.py +6 -7
- waldiez/exporting/tools/exporter.py +5 -0
- waldiez/exporting/tools/factory.py +4 -0
- waldiez/exporting/tools/processor.py +5 -1
- waldiez/io/_ws.py +13 -5
- waldiez/io/models/content/image.py +1 -0
- waldiez/io/models/user_input.py +4 -4
- waldiez/io/models/user_response.py +1 -0
- waldiez/io/mqtt.py +1 -1
- waldiez/io/structured.py +17 -17
- waldiez/io/utils.py +1 -1
- waldiez/io/ws.py +9 -11
- waldiez/logger.py +180 -63
- waldiez/models/agents/agent/update_system_message.py +0 -2
- waldiez/models/agents/doc_agent/doc_agent.py +8 -1
- waldiez/models/common/dict_utils.py +169 -40
- waldiez/models/flow/flow.py +6 -6
- waldiez/models/flow/info.py +5 -1
- waldiez/models/model/_llm.py +28 -14
- waldiez/models/model/model.py +4 -1
- waldiez/models/model/model_data.py +18 -5
- waldiez/models/tool/predefined/_config.py +5 -1
- waldiez/models/tool/predefined/_duckduckgo.py +4 -0
- waldiez/models/tool/predefined/_email.py +474 -0
- waldiez/models/tool/predefined/_google.py +8 -6
- waldiez/models/tool/predefined/_perplexity.py +3 -0
- waldiez/models/tool/predefined/_searxng.py +3 -0
- waldiez/models/tool/predefined/_tavily.py +4 -1
- waldiez/models/tool/predefined/_wikipedia.py +4 -1
- waldiez/models/tool/predefined/_youtube.py +4 -1
- waldiez/models/tool/predefined/protocol.py +3 -0
- waldiez/models/tool/tool.py +22 -4
- waldiez/models/waldiez.py +12 -0
- waldiez/runner.py +37 -54
- waldiez/running/__init__.py +6 -0
- waldiez/running/base_runner.py +310 -353
- waldiez/running/environment.py +1 -0
- waldiez/running/exceptions.py +9 -0
- waldiez/running/post_run.py +4 -4
- waldiez/running/pre_run.py +51 -40
- waldiez/running/protocol.py +21 -101
- waldiez/running/run_results.py +1 -1
- waldiez/running/standard_runner.py +84 -277
- waldiez/running/step_by_step/__init__.py +46 -0
- waldiez/running/step_by_step/breakpoints_mixin.py +188 -0
- waldiez/running/step_by_step/step_by_step_models.py +224 -0
- waldiez/running/step_by_step/step_by_step_runner.py +745 -0
- waldiez/running/subprocess_runner/__base__.py +282 -0
- waldiez/running/subprocess_runner/__init__.py +16 -0
- waldiez/running/subprocess_runner/_async_runner.py +362 -0
- waldiez/running/subprocess_runner/_sync_runner.py +455 -0
- waldiez/running/subprocess_runner/runner.py +561 -0
- waldiez/running/timeline_processor.py +1 -1
- waldiez/running/utils.py +376 -1
- waldiez/utils/version.py +2 -6
- waldiez/ws/__init__.py +70 -0
- waldiez/ws/__main__.py +15 -0
- waldiez/ws/_file_handler.py +201 -0
- waldiez/ws/cli.py +211 -0
- waldiez/ws/client_manager.py +835 -0
- waldiez/ws/errors.py +416 -0
- waldiez/ws/models.py +971 -0
- waldiez/ws/reloader.py +342 -0
- waldiez/ws/server.py +469 -0
- waldiez/ws/session_manager.py +393 -0
- waldiez/ws/session_stats.py +83 -0
- waldiez/ws/utils.py +385 -0
- {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/METADATA +74 -74
- {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/RECORD +87 -65
- waldiez/running/patch_io_stream.py +0 -210
- {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/WHEEL +0 -0
- {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/entry_points.txt +0 -0
- {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/licenses/LICENSE +0 -0
- {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/licenses/NOTICE.md +0 -0
|
@@ -0,0 +1,835 @@
|
|
|
1
|
+
# SPDX-License-Identifier: Apache-2.0.
|
|
2
|
+
# Copyright (c) 2024 - 2025 Waldiez and contributors.
|
|
3
|
+
# pylint: disable=too-many-try-statements,broad-exception-caught
|
|
4
|
+
# pylint: disable=too-complex,too-many-return-statements
|
|
5
|
+
# pyright: reportUnknownMemberType=false,reportAttributeAccessIssue=false
|
|
6
|
+
# pyright: reportUnknownVariableType=false,reportUnknownArgumentType=false
|
|
7
|
+
# pyright: reportAssignmentType=false
|
|
8
|
+
# flake8: noqa: C901
|
|
9
|
+
"""WebSocket client manager: bridges WS <-> subprocess runner."""
|
|
10
|
+
|
|
11
|
+
import asyncio
|
|
12
|
+
import json
|
|
13
|
+
import logging
|
|
14
|
+
import time
|
|
15
|
+
from pathlib import Path
|
|
16
|
+
from typing import Any, Callable, Literal
|
|
17
|
+
|
|
18
|
+
import websockets
|
|
19
|
+
|
|
20
|
+
from waldiez.models import Waldiez
|
|
21
|
+
from waldiez.running.subprocess_runner.runner import WaldiezSubprocessRunner
|
|
22
|
+
|
|
23
|
+
from ._file_handler import FileRequestHandler
|
|
24
|
+
from .errors import (
|
|
25
|
+
ErrorHandler,
|
|
26
|
+
MessageParsingError,
|
|
27
|
+
NoInputRequestedError,
|
|
28
|
+
SessionNotFoundError,
|
|
29
|
+
StaleInputRequestError,
|
|
30
|
+
UnsupportedActionError,
|
|
31
|
+
)
|
|
32
|
+
from .models import (
|
|
33
|
+
BreakpointRequest,
|
|
34
|
+
BreakpointResponse,
|
|
35
|
+
ConvertWorkflowRequest,
|
|
36
|
+
ExecutionMode,
|
|
37
|
+
GetStatusRequest,
|
|
38
|
+
PingRequest,
|
|
39
|
+
PongResponse,
|
|
40
|
+
RunWorkflowRequest,
|
|
41
|
+
RunWorkflowResponse,
|
|
42
|
+
SaveFlowRequest,
|
|
43
|
+
StatusResponse,
|
|
44
|
+
StepControlRequest,
|
|
45
|
+
StepControlResponse,
|
|
46
|
+
StepDebugNotification,
|
|
47
|
+
StepRunWorkflowRequest,
|
|
48
|
+
StepRunWorkflowResponse,
|
|
49
|
+
SubprocessCompletionNotification,
|
|
50
|
+
SubprocessOutputNotification,
|
|
51
|
+
UserInputRequestNotification,
|
|
52
|
+
UserInputResponse,
|
|
53
|
+
WorkflowCompletionNotification,
|
|
54
|
+
WorkflowStatus,
|
|
55
|
+
WorkflowStatusNotification,
|
|
56
|
+
create_error_response,
|
|
57
|
+
parse_client_message,
|
|
58
|
+
)
|
|
59
|
+
from .session_manager import SessionManager
|
|
60
|
+
|
|
61
|
+
CWD = Path.cwd()
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
# pylint: disable=too-many-instance-attributes
|
|
65
|
+
class ClientManager:
|
|
66
|
+
"""Single websocket client and route messages to subprocess runners."""
|
|
67
|
+
|
|
68
|
+
def __init__(
|
|
69
|
+
self,
|
|
70
|
+
websocket: websockets.ServerConnection,
|
|
71
|
+
client_id: str,
|
|
72
|
+
session_manager: SessionManager,
|
|
73
|
+
workspace_dir: Path = CWD,
|
|
74
|
+
error_handler: ErrorHandler | None = None,
|
|
75
|
+
) -> None:
|
|
76
|
+
self.websocket = websocket
|
|
77
|
+
self.client_id = client_id
|
|
78
|
+
self.session_manager = session_manager
|
|
79
|
+
self.workspace_dir = workspace_dir
|
|
80
|
+
self.is_active = True
|
|
81
|
+
|
|
82
|
+
# Active runners per session
|
|
83
|
+
self._runners: dict[str, WaldiezSubprocessRunner] = {}
|
|
84
|
+
|
|
85
|
+
# Track pending input requests (session_id -> last request_id)
|
|
86
|
+
self._pending_input: dict[str, str] = {}
|
|
87
|
+
self._last_prompt: dict[str, str] = {}
|
|
88
|
+
|
|
89
|
+
self.connection_time = time.time()
|
|
90
|
+
|
|
91
|
+
# Extract client info
|
|
92
|
+
self.remote_address = websocket.remote_address
|
|
93
|
+
if not websocket.request:
|
|
94
|
+
raise ValueError("WebSocket request is not available")
|
|
95
|
+
self.user_agent = websocket.request.headers.get("User-Agent", "Unknown")
|
|
96
|
+
# getattr(websocket, "request_headers", {}).get("User-Agent", "Unknown")
|
|
97
|
+
self.logger = logging.getLogger(__name__)
|
|
98
|
+
self.logger.info(
|
|
99
|
+
"Client connected: %s from %s (User-Agent: %s)",
|
|
100
|
+
self.client_id,
|
|
101
|
+
self.remote_address,
|
|
102
|
+
self.user_agent,
|
|
103
|
+
)
|
|
104
|
+
self.error_handler = error_handler or ErrorHandler(self.logger)
|
|
105
|
+
|
|
106
|
+
self.loop: asyncio.AbstractEventLoop | None = None
|
|
107
|
+
try:
|
|
108
|
+
self.loop = asyncio.get_running_loop()
|
|
109
|
+
except RuntimeError:
|
|
110
|
+
# constructed from a sync context? (tests?)
|
|
111
|
+
self.loop = None
|
|
112
|
+
|
|
113
|
+
@property
|
|
114
|
+
def connection_duration(self) -> float:
|
|
115
|
+
"""Get connection duration in seconds."""
|
|
116
|
+
return time.time() - self.connection_time
|
|
117
|
+
|
|
118
|
+
# ---------------------------------------------------------------------
|
|
119
|
+
# Outbound (server -> client)
|
|
120
|
+
# ---------------------------------------------------------------------
|
|
121
|
+
|
|
122
|
+
async def send_message(self, payload: dict[str, Any] | Any) -> bool:
|
|
123
|
+
"""Serialize and send a message to the client.
|
|
124
|
+
|
|
125
|
+
Parameters
|
|
126
|
+
----------
|
|
127
|
+
payload : dict[str, Any] | Any
|
|
128
|
+
The message payload to send.
|
|
129
|
+
|
|
130
|
+
Returns
|
|
131
|
+
-------
|
|
132
|
+
bool
|
|
133
|
+
True if the message was sent successfully, False otherwise.
|
|
134
|
+
"""
|
|
135
|
+
try:
|
|
136
|
+
if hasattr(payload, "model_dump"):
|
|
137
|
+
data = payload.model_dump(mode="json", exclude_none=True)
|
|
138
|
+
elif isinstance(payload, dict):
|
|
139
|
+
data = payload
|
|
140
|
+
else:
|
|
141
|
+
data = json.loads(json.dumps(payload, default=str))
|
|
142
|
+
await self.websocket.send(json.dumps(data))
|
|
143
|
+
return True
|
|
144
|
+
except (websockets.ConnectionClosed, ConnectionResetError) as e:
|
|
145
|
+
self.logger.info("Client %s disconnected: %s", self.client_id, e)
|
|
146
|
+
await self.cleanup()
|
|
147
|
+
return False
|
|
148
|
+
except Exception as e: # pragma: no cover
|
|
149
|
+
self.logger.warning(
|
|
150
|
+
"Failed sending to client %s: %s", self.client_id, e
|
|
151
|
+
)
|
|
152
|
+
# Record operational error
|
|
153
|
+
self.error_handler.record_send_failure(self.client_id)
|
|
154
|
+
return False
|
|
155
|
+
|
|
156
|
+
def close_connection(self) -> None:
|
|
157
|
+
"""Mark as inactive (server will close the socket elsewhere)."""
|
|
158
|
+
self.is_active = False
|
|
159
|
+
|
|
160
|
+
async def cleanup(self) -> None:
|
|
161
|
+
"""Clean up resources when client disconnects."""
|
|
162
|
+
for session_id, runner in self._runners.items():
|
|
163
|
+
try:
|
|
164
|
+
runner.stop()
|
|
165
|
+
await self.session_manager.remove_session(session_id)
|
|
166
|
+
except Exception as e:
|
|
167
|
+
self.logger.warning(
|
|
168
|
+
"Error cleaning up session %s: %s ", session_id, e
|
|
169
|
+
)
|
|
170
|
+
|
|
171
|
+
self._runners.clear()
|
|
172
|
+
self._pending_input.clear()
|
|
173
|
+
self._last_prompt.clear()
|
|
174
|
+
self.close_connection()
|
|
175
|
+
|
|
176
|
+
# ---------------------------------------------------------------------
|
|
177
|
+
# Runner -> Client bridges
|
|
178
|
+
# ---------------------------------------------------------------------
|
|
179
|
+
|
|
180
|
+
def _mk_on_output(
|
|
181
|
+
self, session_id: str
|
|
182
|
+
) -> Callable[[dict[str, Any]], None]:
|
|
183
|
+
"""Runner output callback (called from runner threads)."""
|
|
184
|
+
|
|
185
|
+
def _cb(data: dict[str, Any]) -> None:
|
|
186
|
+
loop = self._ensure_loop()
|
|
187
|
+
if loop is None:
|
|
188
|
+
# No running loop available; best we can do is log and drop.
|
|
189
|
+
self.logger.debug(
|
|
190
|
+
"No running loop to post runner output; dropping."
|
|
191
|
+
)
|
|
192
|
+
return
|
|
193
|
+
data = {**data, "session_id": session_id}
|
|
194
|
+
asyncio.run_coroutine_threadsafe(
|
|
195
|
+
self._handle_runner_output(data), loop
|
|
196
|
+
)
|
|
197
|
+
|
|
198
|
+
return _cb
|
|
199
|
+
|
|
200
|
+
# pylint: disable=no-self-use
|
|
201
|
+
def _mk_on_input_request(self, session_id: str) -> Callable[[str], None]:
|
|
202
|
+
"""Runner input-request callback (fallback if prompt-only)."""
|
|
203
|
+
|
|
204
|
+
def _cb(prompt: str) -> None:
|
|
205
|
+
loop = self._ensure_loop()
|
|
206
|
+
if loop is None:
|
|
207
|
+
# No running loop available; best we can do is log and drop.
|
|
208
|
+
self.logger.debug(
|
|
209
|
+
"No running loop to post runner output; dropping."
|
|
210
|
+
)
|
|
211
|
+
return
|
|
212
|
+
|
|
213
|
+
request_id = f"req_{time.monotonic_ns()}"
|
|
214
|
+
self._pending_input[session_id] = request_id
|
|
215
|
+
self._last_prompt[session_id] = prompt or "> "
|
|
216
|
+
|
|
217
|
+
async def notify() -> None:
|
|
218
|
+
try:
|
|
219
|
+
await self.session_manager.update_session_status(
|
|
220
|
+
session_id, WorkflowStatus.INPUT_WAITING
|
|
221
|
+
)
|
|
222
|
+
await self.send_message(
|
|
223
|
+
UserInputRequestNotification(
|
|
224
|
+
session_id=session_id,
|
|
225
|
+
request_id=request_id,
|
|
226
|
+
prompt=prompt or "> ",
|
|
227
|
+
password=False,
|
|
228
|
+
timeout=120.0,
|
|
229
|
+
)
|
|
230
|
+
)
|
|
231
|
+
except Exception as e: # pragma: no cover
|
|
232
|
+
self.logger.warning("Failed to notify input request: %s", e)
|
|
233
|
+
|
|
234
|
+
# hand off to the loop from runner thread
|
|
235
|
+
asyncio.run_coroutine_threadsafe(notify(), loop)
|
|
236
|
+
|
|
237
|
+
return _cb
|
|
238
|
+
|
|
239
|
+
# ---------------------------------------------------------------------
|
|
240
|
+
# Outbound (server -> client)
|
|
241
|
+
# ---------------------------------------------------------------------
|
|
242
|
+
|
|
243
|
+
# pylint: disable=too-many-branches
|
|
244
|
+
async def handle_message(self, raw_message: str) -> dict[str, Any] | None:
|
|
245
|
+
"""Parse & dispatch an inbound client message.
|
|
246
|
+
|
|
247
|
+
Return an immediate *response* dict (serialized later by server),
|
|
248
|
+
or None if we've already sent notifications.
|
|
249
|
+
|
|
250
|
+
Parameters
|
|
251
|
+
----------
|
|
252
|
+
raw_message : str
|
|
253
|
+
The raw message received from the client.
|
|
254
|
+
|
|
255
|
+
Returns
|
|
256
|
+
-------
|
|
257
|
+
dict[str, Any] | None
|
|
258
|
+
The parsed message or None if it couldn't be parsed.
|
|
259
|
+
"""
|
|
260
|
+
try:
|
|
261
|
+
msg = parse_client_message(raw_message)
|
|
262
|
+
except ValueError as e:
|
|
263
|
+
# Wrap in domain error and format consistently
|
|
264
|
+
return self._error_to_response(MessageParsingError(str(e)))
|
|
265
|
+
|
|
266
|
+
# Lightweight utility requests
|
|
267
|
+
if isinstance(msg, PingRequest):
|
|
268
|
+
return PongResponse.ok(echo_data=msg.echo_data).model_dump(
|
|
269
|
+
mode="json"
|
|
270
|
+
)
|
|
271
|
+
|
|
272
|
+
if isinstance(msg, GetStatusRequest):
|
|
273
|
+
server_status = await self.session_manager.get_status()
|
|
274
|
+
wf_status = None
|
|
275
|
+
if msg.session_id:
|
|
276
|
+
sess = await self.session_manager.get_session(msg.session_id)
|
|
277
|
+
wf_status = sess.status if sess else None
|
|
278
|
+
return StatusResponse.ok(
|
|
279
|
+
server_status=server_status,
|
|
280
|
+
workflow_status=wf_status,
|
|
281
|
+
session_id=msg.session_id,
|
|
282
|
+
).model_dump(mode="json")
|
|
283
|
+
|
|
284
|
+
if isinstance(msg, SaveFlowRequest):
|
|
285
|
+
return FileRequestHandler.handle_save_flow_request(
|
|
286
|
+
msg=msg,
|
|
287
|
+
workspace_dir=self.workspace_dir,
|
|
288
|
+
client_id=self.client_id,
|
|
289
|
+
logger=self.logger,
|
|
290
|
+
)
|
|
291
|
+
|
|
292
|
+
if isinstance(msg, ConvertWorkflowRequest):
|
|
293
|
+
return FileRequestHandler.handle_convert_workflow_request(
|
|
294
|
+
msg=msg,
|
|
295
|
+
client_id=self.client_id,
|
|
296
|
+
workspace_dir=self.workspace_dir,
|
|
297
|
+
logger=self.logger,
|
|
298
|
+
)
|
|
299
|
+
|
|
300
|
+
# Start workflow (STANDARD)
|
|
301
|
+
if isinstance(msg, RunWorkflowRequest):
|
|
302
|
+
return await self._handle_run_workflow(msg)
|
|
303
|
+
|
|
304
|
+
# Start workflow (STEP-BY-STEP / DEBUG)
|
|
305
|
+
if isinstance(msg, StepRunWorkflowRequest):
|
|
306
|
+
return await self._handle_step_run_workflow(msg)
|
|
307
|
+
|
|
308
|
+
# Step controls
|
|
309
|
+
if isinstance(msg, StepControlRequest):
|
|
310
|
+
return await self._handle_step_control(msg)
|
|
311
|
+
|
|
312
|
+
# Breakpoint controls
|
|
313
|
+
if isinstance(msg, BreakpointRequest):
|
|
314
|
+
return await self._handle_breakpoint_control(msg)
|
|
315
|
+
|
|
316
|
+
# User input for pending input_request
|
|
317
|
+
if isinstance(msg, UserInputResponse):
|
|
318
|
+
return await self._handle_user_input(msg)
|
|
319
|
+
|
|
320
|
+
# Stop workflow
|
|
321
|
+
if hasattr(msg, "type") and msg.type == "stop_workflow":
|
|
322
|
+
return await self._handle_stop_workflow(msg)
|
|
323
|
+
|
|
324
|
+
# Unknown
|
|
325
|
+
return self._error_to_response(
|
|
326
|
+
UnsupportedActionError(getattr(msg, "type", "unknown"))
|
|
327
|
+
)
|
|
328
|
+
|
|
329
|
+
async def _handle_run_workflow(
|
|
330
|
+
self, msg: RunWorkflowRequest
|
|
331
|
+
) -> dict[str, Any]:
|
|
332
|
+
try:
|
|
333
|
+
data_dict = json.loads(msg.flow_data)
|
|
334
|
+
waldiez = Waldiez.from_dict(data_dict)
|
|
335
|
+
except Exception as e:
|
|
336
|
+
return RunWorkflowResponse.fail(
|
|
337
|
+
error=f"Invalid flow_data: {e}",
|
|
338
|
+
execution_mode=msg.execution_mode,
|
|
339
|
+
session_id="",
|
|
340
|
+
).model_dump(mode="json")
|
|
341
|
+
# structured path preferred
|
|
342
|
+
session_id = self._next_session_id()
|
|
343
|
+
runner = WaldiezSubprocessRunner(
|
|
344
|
+
waldiez=waldiez,
|
|
345
|
+
on_output=self._mk_on_output(session_id),
|
|
346
|
+
on_input_request=self._mk_on_input_request(session_id),
|
|
347
|
+
mode="run",
|
|
348
|
+
)
|
|
349
|
+
|
|
350
|
+
await self._create_session_for_runner(
|
|
351
|
+
runner, ExecutionMode.STANDARD, session_id=session_id
|
|
352
|
+
)
|
|
353
|
+
# Fire and forget running in a thread
|
|
354
|
+
asyncio.create_task(self._run_runner(session_id, runner))
|
|
355
|
+
|
|
356
|
+
return RunWorkflowResponse.ok(
|
|
357
|
+
session_id=session_id, execution_mode=ExecutionMode.STANDARD
|
|
358
|
+
).model_dump(mode="json")
|
|
359
|
+
|
|
360
|
+
async def _handle_step_run_workflow(
|
|
361
|
+
self, msg: StepRunWorkflowRequest
|
|
362
|
+
) -> dict[str, Any]:
|
|
363
|
+
try:
|
|
364
|
+
data_dict = json.loads(msg.flow_data)
|
|
365
|
+
waldiez = Waldiez.from_dict(data_dict)
|
|
366
|
+
except Exception as e:
|
|
367
|
+
return StepRunWorkflowResponse.fail(
|
|
368
|
+
error=f"Invalid flow_data: {e}",
|
|
369
|
+
session_id="",
|
|
370
|
+
auto_continue=msg.auto_continue,
|
|
371
|
+
breakpoints=msg.breakpoints,
|
|
372
|
+
).model_dump(mode="json")
|
|
373
|
+
session_id = self._next_session_id()
|
|
374
|
+
runner = WaldiezSubprocessRunner(
|
|
375
|
+
waldiez=waldiez,
|
|
376
|
+
on_output=self._mk_on_output(session_id),
|
|
377
|
+
on_input_request=self._mk_on_input_request(session_id),
|
|
378
|
+
mode="debug", # step-by-step via CLI
|
|
379
|
+
)
|
|
380
|
+
|
|
381
|
+
await self._create_session_for_runner(
|
|
382
|
+
runner,
|
|
383
|
+
ExecutionMode.STEP_BY_STEP,
|
|
384
|
+
session_id=session_id,
|
|
385
|
+
)
|
|
386
|
+
|
|
387
|
+
# Persist desired flags (optional)
|
|
388
|
+
sess = await self.session_manager.get_session(session_id)
|
|
389
|
+
if sess:
|
|
390
|
+
sess.state.metadata.update(
|
|
391
|
+
{
|
|
392
|
+
"auto_continue": msg.auto_continue,
|
|
393
|
+
"breakpoints": list(msg.breakpoints),
|
|
394
|
+
}
|
|
395
|
+
)
|
|
396
|
+
|
|
397
|
+
asyncio.create_task(self._run_runner(session_id, runner))
|
|
398
|
+
|
|
399
|
+
return StepRunWorkflowResponse.ok(
|
|
400
|
+
session_id=session_id,
|
|
401
|
+
auto_continue=msg.auto_continue,
|
|
402
|
+
breakpoints=list(msg.breakpoints),
|
|
403
|
+
).model_dump(mode="json")
|
|
404
|
+
|
|
405
|
+
async def _create_session_for_runner(
|
|
406
|
+
self,
|
|
407
|
+
runner: WaldiezSubprocessRunner,
|
|
408
|
+
mode: ExecutionMode,
|
|
409
|
+
session_id: str,
|
|
410
|
+
) -> None:
|
|
411
|
+
await self.session_manager.create_session(
|
|
412
|
+
session_id=session_id,
|
|
413
|
+
client_id=self.client_id,
|
|
414
|
+
execution_mode=mode,
|
|
415
|
+
runner=runner,
|
|
416
|
+
temp_file=None,
|
|
417
|
+
metadata={},
|
|
418
|
+
)
|
|
419
|
+
self._runners[session_id] = runner
|
|
420
|
+
await self.session_manager.update_session_status(
|
|
421
|
+
session_id, WorkflowStatus.STARTING
|
|
422
|
+
)
|
|
423
|
+
await self.send_message(
|
|
424
|
+
WorkflowStatusNotification.make(
|
|
425
|
+
session_id, WorkflowStatus.STARTING, mode
|
|
426
|
+
)
|
|
427
|
+
)
|
|
428
|
+
|
|
429
|
+
async def _run_runner(
|
|
430
|
+
self, session_id: str, runner: WaldiezSubprocessRunner
|
|
431
|
+
) -> None:
|
|
432
|
+
"""Run the subprocess in a thread.
|
|
433
|
+
|
|
434
|
+
Completion is reported via on_output (completion message)
|
|
435
|
+
and here as a fallback.
|
|
436
|
+
|
|
437
|
+
Parameters
|
|
438
|
+
----------
|
|
439
|
+
session_id : str
|
|
440
|
+
The ID of the session.
|
|
441
|
+
runner : WaldiezSubprocessRunner
|
|
442
|
+
The runner instance to execute.
|
|
443
|
+
"""
|
|
444
|
+
try:
|
|
445
|
+
await self.session_manager.update_session_status(
|
|
446
|
+
session_id, WorkflowStatus.RUNNING
|
|
447
|
+
)
|
|
448
|
+
# Run in thread to avoid blocking loop
|
|
449
|
+
# noinspection PyTypeChecker
|
|
450
|
+
await asyncio.to_thread(runner.run, mode=runner.mode)
|
|
451
|
+
# If the runner emitted a completion dict,
|
|
452
|
+
# _handle_runner_output will forward it.
|
|
453
|
+
except Exception as e: # pragma: no cover
|
|
454
|
+
await self.session_manager.update_session_status(
|
|
455
|
+
session_id, WorkflowStatus.FAILED
|
|
456
|
+
)
|
|
457
|
+
await self.send_message(
|
|
458
|
+
WorkflowCompletionNotification(
|
|
459
|
+
session_id=session_id,
|
|
460
|
+
success=False,
|
|
461
|
+
exit_code=-1,
|
|
462
|
+
error=str(e),
|
|
463
|
+
)
|
|
464
|
+
)
|
|
465
|
+
|
|
466
|
+
async def _handle_step_control(
|
|
467
|
+
self, msg: StepControlRequest
|
|
468
|
+
) -> dict[str, Any]:
|
|
469
|
+
runner = self._runners.get(msg.session_id)
|
|
470
|
+
if not runner:
|
|
471
|
+
return StepControlResponse.fail(
|
|
472
|
+
error="Session not found",
|
|
473
|
+
action=msg.action,
|
|
474
|
+
result="",
|
|
475
|
+
session_id=msg.session_id,
|
|
476
|
+
).model_dump(mode="json")
|
|
477
|
+
|
|
478
|
+
code = {
|
|
479
|
+
"": "c",
|
|
480
|
+
"continue": "c",
|
|
481
|
+
"step": "s",
|
|
482
|
+
"run": "r",
|
|
483
|
+
"quit": "q",
|
|
484
|
+
"info": "i",
|
|
485
|
+
"help": "h",
|
|
486
|
+
"stats": "st",
|
|
487
|
+
}.get(msg.action)
|
|
488
|
+
|
|
489
|
+
if not code:
|
|
490
|
+
return StepControlResponse.fail(
|
|
491
|
+
error=f"Unsupported action: {msg.action}",
|
|
492
|
+
action=msg.action,
|
|
493
|
+
result="",
|
|
494
|
+
session_id=msg.session_id,
|
|
495
|
+
).model_dump(mode="json")
|
|
496
|
+
|
|
497
|
+
runner.provide_user_input(code)
|
|
498
|
+
return StepControlResponse.ok(
|
|
499
|
+
action=msg.action, result="sent", session_id=msg.session_id
|
|
500
|
+
).model_dump(mode="json")
|
|
501
|
+
|
|
502
|
+
async def _handle_breakpoint_control(
|
|
503
|
+
self, msg: BreakpointRequest
|
|
504
|
+
) -> dict[str, Any]:
|
|
505
|
+
runner = self._runners.get(msg.session_id)
|
|
506
|
+
if not runner:
|
|
507
|
+
return BreakpointResponse.fail(
|
|
508
|
+
error="Session not found",
|
|
509
|
+
action=msg.action,
|
|
510
|
+
session_id=msg.session_id,
|
|
511
|
+
).model_dump(mode="json")
|
|
512
|
+
|
|
513
|
+
cmd = {
|
|
514
|
+
"list": "lb",
|
|
515
|
+
"clear": "cb",
|
|
516
|
+
"add": "ab",
|
|
517
|
+
"remove": "rb",
|
|
518
|
+
}[msg.action]
|
|
519
|
+
|
|
520
|
+
# NOTE: If we later add `ab <event>`, send f"{cmd} {msg.event_type}"
|
|
521
|
+
runner.provide_user_input(cmd)
|
|
522
|
+
return BreakpointResponse.ok(
|
|
523
|
+
action=msg.action, session_id=msg.session_id
|
|
524
|
+
).model_dump(mode="json")
|
|
525
|
+
|
|
526
|
+
async def _handle_user_input(
|
|
527
|
+
self, msg: UserInputResponse
|
|
528
|
+
) -> dict[str, Any]:
|
|
529
|
+
runner = self._runners.get(msg.session_id)
|
|
530
|
+
if not runner:
|
|
531
|
+
return self._error_to_response(
|
|
532
|
+
SessionNotFoundError(session_id=msg.session_id)
|
|
533
|
+
)
|
|
534
|
+
|
|
535
|
+
pending = self._pending_input.get(msg.session_id)
|
|
536
|
+
if not pending:
|
|
537
|
+
return self._error_to_response(NoInputRequestedError())
|
|
538
|
+
if pending != msg.request_id:
|
|
539
|
+
self.logger.debug(
|
|
540
|
+
"Client %s: mismatched request_id for %s (got %s expected %s)",
|
|
541
|
+
self.client_id,
|
|
542
|
+
msg.session_id,
|
|
543
|
+
msg.request_id,
|
|
544
|
+
pending,
|
|
545
|
+
)
|
|
546
|
+
return self._error_to_response(
|
|
547
|
+
StaleInputRequestError(
|
|
548
|
+
request_id=msg.request_id, expected_id=pending
|
|
549
|
+
)
|
|
550
|
+
)
|
|
551
|
+
|
|
552
|
+
runner.provide_user_input(msg.data)
|
|
553
|
+
return {"type": "ok", "success": True}
|
|
554
|
+
|
|
555
|
+
async def _handle_stop_workflow(self, msg: Any) -> dict[str, Any]:
|
|
556
|
+
session_id = getattr(msg, "session_id", "")
|
|
557
|
+
runner = self._runners.get(session_id)
|
|
558
|
+
if not runner:
|
|
559
|
+
return self._error_to_response(
|
|
560
|
+
SessionNotFoundError(session_id=session_id)
|
|
561
|
+
)
|
|
562
|
+
|
|
563
|
+
try:
|
|
564
|
+
runner.stop()
|
|
565
|
+
await self.session_manager.update_session_status(
|
|
566
|
+
session_id, WorkflowStatus.STOPPING
|
|
567
|
+
)
|
|
568
|
+
return {
|
|
569
|
+
"type": "stop_workflow_response",
|
|
570
|
+
"session_id": session_id,
|
|
571
|
+
"success": True,
|
|
572
|
+
"forced": getattr(msg, "force", False),
|
|
573
|
+
}
|
|
574
|
+
except Exception as e:
|
|
575
|
+
# Shape it as a standard error response
|
|
576
|
+
return self._error_to_response(e)
|
|
577
|
+
|
|
578
|
+
async def _handle_runner_output(self, data: dict[str, Any]) -> None:
|
|
579
|
+
"""Handle output from the runner.
|
|
580
|
+
|
|
581
|
+
Parameters
|
|
582
|
+
----------
|
|
583
|
+
data : dict[str, Any[
|
|
584
|
+
The output dict from the runner.
|
|
585
|
+
|
|
586
|
+
Interpret dicts from BaseSubprocessRunner.create_* and parse_output().
|
|
587
|
+
We support:
|
|
588
|
+
- type == "subprocess_output"
|
|
589
|
+
-> SubprocessOutputNotification
|
|
590
|
+
- type == "subprocess_completion"
|
|
591
|
+
-> SubprocessCompletionNotification + status update
|
|
592
|
+
- type in {"input_request", "debug_input_request"}
|
|
593
|
+
-> UserInputRequestNotification
|
|
594
|
+
- type startswith "debug_"
|
|
595
|
+
-> StepDebugNotification
|
|
596
|
+
- anything else
|
|
597
|
+
-> fallback stdout SubprocessOutputNotification
|
|
598
|
+
"""
|
|
599
|
+
try:
|
|
600
|
+
msg_type = str(data.get("type", "")).lower()
|
|
601
|
+
session_id_raw = data.get("session_id")
|
|
602
|
+
session_id = (
|
|
603
|
+
str(session_id_raw)
|
|
604
|
+
if session_id_raw
|
|
605
|
+
else (self._guess_session_id() or "")
|
|
606
|
+
)
|
|
607
|
+
|
|
608
|
+
if msg_type in ("input_request", "debug_input_request"):
|
|
609
|
+
await self._handle_runner_input_request(session_id, data)
|
|
610
|
+
return
|
|
611
|
+
|
|
612
|
+
if msg_type.startswith("debug_"):
|
|
613
|
+
await self._handle_runner_debug(session_id, msg_type, data)
|
|
614
|
+
return
|
|
615
|
+
|
|
616
|
+
if msg_type == "subprocess_completion":
|
|
617
|
+
await self._handle_runner_completion(session_id, data)
|
|
618
|
+
return
|
|
619
|
+
|
|
620
|
+
if msg_type == "subprocess_output":
|
|
621
|
+
await self._handle_runner_subprocess_output(session_id, data)
|
|
622
|
+
return
|
|
623
|
+
|
|
624
|
+
# Fallback: dump everything as stdout line
|
|
625
|
+
await self.send_message(
|
|
626
|
+
SubprocessOutputNotification(
|
|
627
|
+
session_id=session_id,
|
|
628
|
+
stream="stdout",
|
|
629
|
+
content=json.dumps(data, default=str),
|
|
630
|
+
subprocess_type="output",
|
|
631
|
+
context={},
|
|
632
|
+
)
|
|
633
|
+
)
|
|
634
|
+
except Exception as e: # pragma: no cover
|
|
635
|
+
# Convert to a debug notification so the client sees something,
|
|
636
|
+
# and keep the server loop healthy.
|
|
637
|
+
await self.send_message(
|
|
638
|
+
StepDebugNotification(
|
|
639
|
+
session_id=self._guess_session_id() or "",
|
|
640
|
+
debug_type="error",
|
|
641
|
+
data={
|
|
642
|
+
"message": "Runner output handling failed",
|
|
643
|
+
"error": str(e),
|
|
644
|
+
},
|
|
645
|
+
)
|
|
646
|
+
)
|
|
647
|
+
|
|
648
|
+
async def _handle_runner_input_request(
|
|
649
|
+
self, session_id: str, data: dict[str, Any]
|
|
650
|
+
) -> None:
|
|
651
|
+
"""Handle an input request from the runner."""
|
|
652
|
+
request_id = str(data.get("request_id", ""))
|
|
653
|
+
prompt = str(data.get("prompt", "> "))
|
|
654
|
+
if session_id and request_id:
|
|
655
|
+
self._pending_input[session_id] = request_id
|
|
656
|
+
self._last_prompt[session_id] = prompt
|
|
657
|
+
await self.session_manager.update_session_status(
|
|
658
|
+
session_id, WorkflowStatus.INPUT_WAITING
|
|
659
|
+
)
|
|
660
|
+
await self.send_message(
|
|
661
|
+
UserInputRequestNotification(
|
|
662
|
+
session_id=session_id,
|
|
663
|
+
request_id=request_id,
|
|
664
|
+
prompt=prompt,
|
|
665
|
+
password=bool(data.get("password", False)),
|
|
666
|
+
timeout=float(data.get("timeout", 120.0)),
|
|
667
|
+
)
|
|
668
|
+
)
|
|
669
|
+
|
|
670
|
+
# pylint: disable=line-too-long
|
|
671
|
+
async def _handle_runner_debug(
|
|
672
|
+
self, session_id: str, msg_type: str, data: dict[str, Any]
|
|
673
|
+
) -> None:
|
|
674
|
+
"""Handle a debug message from the runner."""
|
|
675
|
+
# noinspection PyTypeChecker
|
|
676
|
+
kind = msg_type.replace("debug_", "", 1) or "info"
|
|
677
|
+
debug_type = (
|
|
678
|
+
kind if kind in {"stats", "help", "error", "info"} else "info"
|
|
679
|
+
) # noqa: E501
|
|
680
|
+
await self.send_message(
|
|
681
|
+
StepDebugNotification(
|
|
682
|
+
session_id=session_id,
|
|
683
|
+
debug_type=debug_type, # type: ignore
|
|
684
|
+
data={
|
|
685
|
+
k: v
|
|
686
|
+
for k, v in data.items()
|
|
687
|
+
if k not in {"type", "session_id"}
|
|
688
|
+
},
|
|
689
|
+
)
|
|
690
|
+
)
|
|
691
|
+
|
|
692
|
+
async def _handle_runner_completion(
|
|
693
|
+
self, session_id: str, data: dict[str, Any]
|
|
694
|
+
) -> None:
|
|
695
|
+
"""Handle a completion message from the runner."""
|
|
696
|
+
success = bool(data.get("success", False))
|
|
697
|
+
exit_code = int(data.get("exit_code", -1))
|
|
698
|
+
message = str(data.get("message", ""))
|
|
699
|
+
|
|
700
|
+
await self.session_manager.update_session_status(
|
|
701
|
+
session_id,
|
|
702
|
+
(WorkflowStatus.COMPLETED if success else WorkflowStatus.FAILED),
|
|
703
|
+
)
|
|
704
|
+
await self.send_message(
|
|
705
|
+
SubprocessCompletionNotification(
|
|
706
|
+
session_id=session_id,
|
|
707
|
+
success=success,
|
|
708
|
+
exit_code=exit_code,
|
|
709
|
+
message=message,
|
|
710
|
+
context=data.get("context", {}) or {},
|
|
711
|
+
)
|
|
712
|
+
)
|
|
713
|
+
|
|
714
|
+
async def _handle_runner_subprocess_output(
|
|
715
|
+
self, session_id: str, data: dict[str, Any]
|
|
716
|
+
) -> None:
|
|
717
|
+
"""Handle a subprocess output message from the runner."""
|
|
718
|
+
stream_ = str(data.get("stream", "stdout"))
|
|
719
|
+
stream: Literal["stdout", "stderr"] = (
|
|
720
|
+
"stderr" if stream_ == "stderr" else "stdout"
|
|
721
|
+
)
|
|
722
|
+
subprocess_type_str = str(data.get("subprocess_type", "output")).lower()
|
|
723
|
+
subprocess_type: Literal["output", "error", "debug"] = (
|
|
724
|
+
subprocess_type_str # type: ignore
|
|
725
|
+
if subprocess_type_str in ("output", "error", "debug")
|
|
726
|
+
else "output"
|
|
727
|
+
)
|
|
728
|
+
content = data.get("content", "")
|
|
729
|
+
# noinspection PyUnreachableCode
|
|
730
|
+
if not isinstance(content, str):
|
|
731
|
+
# noinspection PyBroadException
|
|
732
|
+
try:
|
|
733
|
+
content = json.dumps(content, default=str)
|
|
734
|
+
except Exception: # pragma: no cover
|
|
735
|
+
content = str(content)
|
|
736
|
+
context = data.get("context", {}) or {}
|
|
737
|
+
content = self._strip_prompt_if_needed(
|
|
738
|
+
session_id=session_id,
|
|
739
|
+
content=content,
|
|
740
|
+
)
|
|
741
|
+
await self.send_message(
|
|
742
|
+
SubprocessOutputNotification(
|
|
743
|
+
session_id=session_id,
|
|
744
|
+
stream=stream,
|
|
745
|
+
content=content,
|
|
746
|
+
subprocess_type=subprocess_type,
|
|
747
|
+
context=context,
|
|
748
|
+
)
|
|
749
|
+
)
|
|
750
|
+
|
|
751
|
+
def _strip_prompt_if_needed(
|
|
752
|
+
self,
|
|
753
|
+
session_id: str,
|
|
754
|
+
content: str,
|
|
755
|
+
) -> str:
|
|
756
|
+
# let's try to avoid messages like (including the prompt):
|
|
757
|
+
# \"content\": \"Prompt msg: {\\\"type\\\": \\\"text\\\",...\\}\"
|
|
758
|
+
last = self._last_prompt.get(session_id)
|
|
759
|
+
if last:
|
|
760
|
+
s_last = last.strip()
|
|
761
|
+
s_cont = content.strip()
|
|
762
|
+
if s_cont.startswith(s_last):
|
|
763
|
+
self._last_prompt.pop(session_id, None)
|
|
764
|
+
# Remove the leading prompt text
|
|
765
|
+
# Prefer removing from the original string to preserve
|
|
766
|
+
# spacing after the prefix.
|
|
767
|
+
idx = content.find(last)
|
|
768
|
+
payload = content[idx + len(last) :].lstrip()
|
|
769
|
+
|
|
770
|
+
# If the remainder looks like JSON,
|
|
771
|
+
# try to parse and re-dispatch it
|
|
772
|
+
if (
|
|
773
|
+
payload
|
|
774
|
+
and payload[0] in ("{", "[")
|
|
775
|
+
and payload[-1] in ("}", "]")
|
|
776
|
+
):
|
|
777
|
+
# noinspection TryExceptPass,PyBroadException
|
|
778
|
+
try:
|
|
779
|
+
json.loads(payload)
|
|
780
|
+
except Exception:
|
|
781
|
+
# Not valid JSON—fall through
|
|
782
|
+
# just forward the cleaned text
|
|
783
|
+
pass
|
|
784
|
+
else:
|
|
785
|
+
return payload
|
|
786
|
+
|
|
787
|
+
# If not JSON, keep the cleaned textual payload
|
|
788
|
+
return payload
|
|
789
|
+
return content
|
|
790
|
+
|
|
791
|
+
def _guess_session_id(self) -> str | None:
|
|
792
|
+
"""Pick any current session for this client (we typically run one)."""
|
|
793
|
+
return next(iter(self._runners.keys()), None)
|
|
794
|
+
|
|
795
|
+
def _next_session_id(self) -> str:
|
|
796
|
+
return f"session_{self.client_id}_{len(self._runners) + 1:02d}"
|
|
797
|
+
|
|
798
|
+
def _ensure_loop(self) -> asyncio.AbstractEventLoop | None:
|
|
799
|
+
"""Return an event loop for scheduling callbacks, if available."""
|
|
800
|
+
if self.loop and not self.loop.is_closed():
|
|
801
|
+
return self.loop
|
|
802
|
+
try:
|
|
803
|
+
# if we're now on an async path, this will work
|
|
804
|
+
self.loop = asyncio.get_running_loop()
|
|
805
|
+
return self.loop
|
|
806
|
+
except RuntimeError:
|
|
807
|
+
return None
|
|
808
|
+
|
|
809
|
+
def _error_to_response(
|
|
810
|
+
self, error: Exception, request_id: str | None = None
|
|
811
|
+
) -> dict[str, Any]:
|
|
812
|
+
"""Convert an exception into a standardized ErrorResponse payload.
|
|
813
|
+
|
|
814
|
+
Parameters
|
|
815
|
+
----------
|
|
816
|
+
error : Exception
|
|
817
|
+
The error to convert
|
|
818
|
+
request_id : str | None
|
|
819
|
+
Optional request id to echo back
|
|
820
|
+
|
|
821
|
+
Returns
|
|
822
|
+
-------
|
|
823
|
+
dict[str, Any]
|
|
824
|
+
Serialized ErrorResponse
|
|
825
|
+
"""
|
|
826
|
+
details = self.error_handler.handle_error(
|
|
827
|
+
error, client_id=self.client_id
|
|
828
|
+
)
|
|
829
|
+
return create_error_response(
|
|
830
|
+
error_message=details.get("message", "An error occurred"),
|
|
831
|
+
error_code=int(details.get("code", 500)),
|
|
832
|
+
error_type=str(details.get("error_type", "InternalError")),
|
|
833
|
+
request_id=request_id,
|
|
834
|
+
details=details.get("details", {}),
|
|
835
|
+
).model_dump(mode="json")
|