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.

Files changed (88) hide show
  1. waldiez/_version.py +1 -1
  2. waldiez/cli.py +112 -24
  3. waldiez/exporting/agent/exporter.py +3 -0
  4. waldiez/exporting/agent/extras/captain_agent_extras.py +44 -7
  5. waldiez/exporting/agent/extras/handoffs/condition.py +3 -1
  6. waldiez/exporting/chats/utils/common.py +25 -23
  7. waldiez/exporting/core/__init__.py +0 -2
  8. waldiez/exporting/core/context.py +13 -13
  9. waldiez/exporting/core/protocols.py +0 -141
  10. waldiez/exporting/core/result.py +5 -5
  11. waldiez/exporting/flow/merger.py +2 -2
  12. waldiez/exporting/flow/orchestrator.py +1 -0
  13. waldiez/exporting/flow/utils/common.py +2 -2
  14. waldiez/exporting/flow/utils/importing.py +1 -0
  15. waldiez/exporting/flow/utils/logging.py +6 -7
  16. waldiez/exporting/tools/exporter.py +5 -0
  17. waldiez/exporting/tools/factory.py +4 -0
  18. waldiez/exporting/tools/processor.py +5 -1
  19. waldiez/io/_ws.py +13 -5
  20. waldiez/io/models/content/image.py +1 -0
  21. waldiez/io/models/user_input.py +4 -4
  22. waldiez/io/models/user_response.py +1 -0
  23. waldiez/io/mqtt.py +1 -1
  24. waldiez/io/structured.py +17 -17
  25. waldiez/io/utils.py +1 -1
  26. waldiez/io/ws.py +9 -11
  27. waldiez/logger.py +180 -63
  28. waldiez/models/agents/agent/update_system_message.py +0 -2
  29. waldiez/models/agents/doc_agent/doc_agent.py +8 -1
  30. waldiez/models/common/dict_utils.py +169 -40
  31. waldiez/models/flow/flow.py +6 -6
  32. waldiez/models/flow/info.py +5 -1
  33. waldiez/models/model/_llm.py +28 -14
  34. waldiez/models/model/model.py +4 -1
  35. waldiez/models/model/model_data.py +18 -5
  36. waldiez/models/tool/predefined/_config.py +5 -1
  37. waldiez/models/tool/predefined/_duckduckgo.py +4 -0
  38. waldiez/models/tool/predefined/_email.py +474 -0
  39. waldiez/models/tool/predefined/_google.py +8 -6
  40. waldiez/models/tool/predefined/_perplexity.py +3 -0
  41. waldiez/models/tool/predefined/_searxng.py +3 -0
  42. waldiez/models/tool/predefined/_tavily.py +4 -1
  43. waldiez/models/tool/predefined/_wikipedia.py +4 -1
  44. waldiez/models/tool/predefined/_youtube.py +4 -1
  45. waldiez/models/tool/predefined/protocol.py +3 -0
  46. waldiez/models/tool/tool.py +22 -4
  47. waldiez/models/waldiez.py +12 -0
  48. waldiez/runner.py +37 -54
  49. waldiez/running/__init__.py +6 -0
  50. waldiez/running/base_runner.py +310 -353
  51. waldiez/running/environment.py +1 -0
  52. waldiez/running/exceptions.py +9 -0
  53. waldiez/running/post_run.py +4 -4
  54. waldiez/running/pre_run.py +51 -40
  55. waldiez/running/protocol.py +21 -101
  56. waldiez/running/run_results.py +1 -1
  57. waldiez/running/standard_runner.py +84 -277
  58. waldiez/running/step_by_step/__init__.py +46 -0
  59. waldiez/running/step_by_step/breakpoints_mixin.py +188 -0
  60. waldiez/running/step_by_step/step_by_step_models.py +224 -0
  61. waldiez/running/step_by_step/step_by_step_runner.py +745 -0
  62. waldiez/running/subprocess_runner/__base__.py +282 -0
  63. waldiez/running/subprocess_runner/__init__.py +16 -0
  64. waldiez/running/subprocess_runner/_async_runner.py +362 -0
  65. waldiez/running/subprocess_runner/_sync_runner.py +455 -0
  66. waldiez/running/subprocess_runner/runner.py +561 -0
  67. waldiez/running/timeline_processor.py +1 -1
  68. waldiez/running/utils.py +376 -1
  69. waldiez/utils/version.py +2 -6
  70. waldiez/ws/__init__.py +70 -0
  71. waldiez/ws/__main__.py +15 -0
  72. waldiez/ws/_file_handler.py +201 -0
  73. waldiez/ws/cli.py +211 -0
  74. waldiez/ws/client_manager.py +835 -0
  75. waldiez/ws/errors.py +416 -0
  76. waldiez/ws/models.py +971 -0
  77. waldiez/ws/reloader.py +342 -0
  78. waldiez/ws/server.py +469 -0
  79. waldiez/ws/session_manager.py +393 -0
  80. waldiez/ws/session_stats.py +83 -0
  81. waldiez/ws/utils.py +385 -0
  82. {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/METADATA +74 -74
  83. {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/RECORD +87 -65
  84. waldiez/running/patch_io_stream.py +0 -210
  85. {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/WHEEL +0 -0
  86. {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/entry_points.txt +0 -0
  87. {waldiez-0.5.8.dist-info → waldiez-0.5.10.dist-info}/licenses/LICENSE +0 -0
  88. {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")