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
waldiez/ws/models.py ADDED
@@ -0,0 +1,971 @@
1
+ # SPDX-License-Identifier: Apache-2.0.
2
+ # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
+ # pylint: disable=broad-exception-caught,no-member
4
+ # pyright: reportUnknownVariableType=false
5
+ """Session management models for WebSocket workflow execution."""
6
+
7
+ import json
8
+ import time
9
+ import uuid
10
+ from enum import Enum
11
+ from pathlib import Path
12
+ from typing import Annotated, Any, Literal, Union
13
+
14
+ from pydantic import BaseModel, ConfigDict, Field, ValidationError
15
+
16
+
17
+ class WorkflowStatus(str, Enum):
18
+ """Workflow execution status."""
19
+
20
+ IDLE = "idle"
21
+ STARTING = "starting"
22
+ RUNNING = "running"
23
+ PAUSED = "paused"
24
+ STEP_WAITING = "step_waiting"
25
+ INPUT_WAITING = "input_waiting"
26
+ STOPPING = "stopping"
27
+ COMPLETED = "completed"
28
+ FAILED = "failed"
29
+ CANCELLED = "cancelled"
30
+
31
+
32
+ class ExecutionMode(str, Enum):
33
+ """Workflow execution modes."""
34
+
35
+ STANDARD = "standard"
36
+ STEP_BY_STEP = "step_by_step"
37
+ SUBPROCESS = "subprocess"
38
+
39
+
40
+ class BaseWaldiezMessage(BaseModel):
41
+ """Base class for all Waldiez WebSocket messages."""
42
+
43
+ message_id: str = Field(
44
+ default_factory=lambda: f"msg_{time.monotonic_ns()}",
45
+ )
46
+ timestamp: float = Field(default_factory=time.time)
47
+
48
+ model_config = ConfigDict(
49
+ arbitrary_types_allowed=True,
50
+ extra="ignore",
51
+ json_encoders={
52
+ Path: str,
53
+ Enum: lambda v: v.value,
54
+ },
55
+ )
56
+
57
+
58
+ class BaseRequest(BaseWaldiezMessage):
59
+ """Base class for client requests."""
60
+
61
+
62
+ class BaseResponse(BaseWaldiezMessage):
63
+ """Base class for server responses."""
64
+
65
+ request_id: str | None = None
66
+ success: bool = True
67
+ error: str | None = None
68
+
69
+ @classmethod
70
+ def ok(cls, **kwargs: Any) -> "BaseResponse":
71
+ """Create a successful response.
72
+
73
+ Parameters
74
+ ----------
75
+ kwargs : Any
76
+ Additional keyword arguments to include in the response.
77
+
78
+ Returns
79
+ -------
80
+ BaseResponse
81
+ The created successful response.
82
+
83
+ """
84
+ kwargs.setdefault("success", True)
85
+ kwargs.setdefault("error", None)
86
+ return cls(**kwargs)
87
+
88
+ @classmethod
89
+ def fail(cls, error: str, **kwargs: Any) -> "BaseResponse":
90
+ """
91
+ Create a failed response.
92
+
93
+ Parameters
94
+ ----------
95
+ error : str
96
+ The error message.
97
+ kwargs : Any
98
+ Additional keyword arguments to include in the response.
99
+
100
+ Returns
101
+ -------
102
+ BaseResponse
103
+ The created failed response.
104
+
105
+ """
106
+ kwargs.update({"success": False, "error": error})
107
+ return cls(**kwargs)
108
+
109
+
110
+ class BaseNotification(BaseWaldiezMessage):
111
+ """Base class for server notifications (no response expected)."""
112
+
113
+
114
+ # ========================================
115
+ # CLIENT-TO-SERVER MESSAGES (REQUESTS)
116
+ # ========================================
117
+
118
+
119
+ class SaveFlowRequest(BaseRequest):
120
+ """Request to save a workflow."""
121
+
122
+ type: Literal["save_flow"] = "save_flow"
123
+ flow_data: str # JSON string of workflow
124
+ filename: str | None = None
125
+ force_overwrite: bool = False
126
+
127
+
128
+ class RunWorkflowRequest(BaseRequest):
129
+ """Request to run a workflow in standard mode."""
130
+
131
+ type: Literal["run_workflow"] = "run_workflow"
132
+ flow_data: str # JSON string of workflow
133
+ execution_mode: ExecutionMode = ExecutionMode.STANDARD
134
+ structured_io: bool = True
135
+ uploads_root: str | None = None
136
+ dot_env_path: str | None = None
137
+ output_path: str | None = None
138
+
139
+
140
+ class StepRunWorkflowRequest(BaseRequest):
141
+ """Request to run a workflow in step-by-step mode."""
142
+
143
+ type: Literal["step_run_workflow"] = "step_run_workflow"
144
+ flow_data: str # JSON string of workflow
145
+ auto_continue: bool = False
146
+ breakpoints: list[str] = Field(default_factory=list)
147
+ structured_io: bool = True
148
+ uploads_root: str | None = None
149
+ dot_env_path: str | None = None
150
+ output_path: str | None = None
151
+
152
+
153
+ class StepControlRequest(BaseRequest):
154
+ """Request to control step-by-step execution."""
155
+
156
+ type: Literal["step_control"] = "step_control"
157
+ action: Literal["continue", "step", "run", "quit", "info", "help", "stats"]
158
+ session_id: str
159
+
160
+
161
+ class BreakpointRequest(BaseRequest):
162
+ """Request to manage breakpoints."""
163
+
164
+ type: Literal["breakpoint_control"] = "breakpoint_control"
165
+ action: Literal["add", "remove", "list", "clear"]
166
+ event_type: str | None = None # Required for add/remove
167
+ session_id: str
168
+
169
+ @property
170
+ def response(self) -> Literal["added", "removed", "list", "cleared"]:
171
+ """Get the response message for the breakpoint action."""
172
+ if self.action == "remove":
173
+ return "removed"
174
+ if self.action == "list":
175
+ return "list"
176
+ if self.action == "clear":
177
+ return "cleared"
178
+ return "added"
179
+
180
+
181
+ class UserInputResponse(BaseRequest):
182
+ """User input response for workflow execution."""
183
+
184
+ type: Literal["user_input"] = "user_input"
185
+ request_id: str
186
+ data: Any
187
+ session_id: str
188
+
189
+
190
+ class StopWorkflowRequest(BaseRequest):
191
+ """Request to stop workflow execution."""
192
+
193
+ type: Literal["stop_workflow"] = "stop_workflow"
194
+ session_id: str
195
+ force: bool = False
196
+
197
+
198
+ class StopWorkflowResponse(BaseResponse):
199
+ """Response to stop workflow request."""
200
+
201
+ session_id: str
202
+ type: Literal["stop_workflow_response"] = "stop_workflow_response"
203
+ error: str | None = None
204
+ forced: bool = False
205
+
206
+
207
+ class ConvertWorkflowRequest(BaseRequest):
208
+ """Request to convert workflow to different format."""
209
+
210
+ type: Literal["convert_workflow"] = "convert_workflow"
211
+ flow_data: str
212
+ target_format: Literal["py", "ipynb"]
213
+ output_path: str | None = None
214
+
215
+
216
+ class UploadFileRequest(BaseRequest):
217
+ """Request to upload a file."""
218
+
219
+ type: Literal["upload_file"] = "upload_file"
220
+ filename: str
221
+ file_data: str # Base64 encoded
222
+ file_size: int
223
+ mime_type: str | None = None
224
+
225
+
226
+ class PingRequest(BaseRequest):
227
+ """Ping request for connection testing."""
228
+
229
+ type: Literal["ping"] = "ping"
230
+ echo_data: dict[str, Any] = Field(default_factory=dict)
231
+
232
+
233
+ class GetStatusRequest(BaseRequest):
234
+ """Request current server/workflow status."""
235
+
236
+ type: Literal["get_status"] = "get_status"
237
+ session_id: str | None = None
238
+
239
+
240
+ # ========================================
241
+ # SERVER-TO-CLIENT MESSAGES (RESPONSES)
242
+ # ========================================
243
+
244
+
245
+ class SaveFlowResponse(BaseResponse):
246
+ """Response to save flow request."""
247
+
248
+ type: Literal["save_flow_response"] = "save_flow_response"
249
+ file_path: str | None = None
250
+ error: str | None = None
251
+
252
+
253
+ class RunWorkflowResponse(BaseResponse):
254
+ """Response to run workflow request."""
255
+
256
+ type: Literal["run_workflow_response"] = "run_workflow_response"
257
+ session_id: str
258
+ execution_mode: ExecutionMode
259
+ error: str | None = None
260
+
261
+
262
+ class StepRunWorkflowResponse(BaseResponse):
263
+ """Response to step run workflow request."""
264
+
265
+ type: Literal["step_run_workflow_response"] = "step_run_workflow_response"
266
+ session_id: str
267
+ auto_continue: bool
268
+ breakpoints: list[str]
269
+ error: str | None = None
270
+
271
+
272
+ class StepControlResponse(BaseResponse):
273
+ """Response to step control request."""
274
+
275
+ type: Literal["step_control_response"] = "step_control_response"
276
+ action: str
277
+ result: str
278
+ session_id: str
279
+
280
+
281
+ class BreakpointResponse(BaseResponse):
282
+ """Response to breakpoint management."""
283
+
284
+ type: Literal["breakpoint_response"] = "breakpoint_response"
285
+ action: str
286
+ breakpoints: list[str] = Field(default_factory=list)
287
+ session_id: str
288
+ error: str | None = None
289
+
290
+
291
+ class ConvertWorkflowResponse(BaseResponse):
292
+ """Response to convert workflow request."""
293
+
294
+ type: Literal["convert_workflow_response"] = "convert_workflow_response"
295
+ target_format: str
296
+ output_path: str | None = None
297
+ error: str | None = None
298
+
299
+
300
+ class UploadFileResponse(BaseResponse):
301
+ """Response to file upload."""
302
+
303
+ type: Literal["upload_file_response"] = "upload_file_response"
304
+ file_path: str | None = None
305
+ file_size: int | None = None
306
+ error: str | None = None
307
+
308
+
309
+ class PongResponse(BaseResponse):
310
+ """Response to ping."""
311
+
312
+ type: Literal["pong"] = "pong"
313
+ echo_data: dict[str, Any] = Field(default_factory=dict)
314
+ server_time: float = Field(default_factory=time.time)
315
+
316
+
317
+ class StatusResponse(BaseResponse):
318
+ """Response with current status."""
319
+
320
+ type: Literal["status_response"] = "status_response"
321
+ server_status: dict[str, Any]
322
+ workflow_status: WorkflowStatus | None = None
323
+ session_id: str | None = None
324
+
325
+
326
+ class ErrorResponse(BaseResponse):
327
+ """Generic error response."""
328
+
329
+ type: Literal["error"] = "error"
330
+ error_code: int
331
+ error_type: str
332
+ details: dict[str, Any] = Field(default_factory=dict)
333
+ success: bool = False
334
+
335
+
336
+ # ========================================
337
+ # SERVER-TO-CLIENT NOTIFICATIONS
338
+ # ========================================
339
+
340
+
341
+ class WorkflowStatusNotification(BaseNotification):
342
+ """Notification of workflow status change."""
343
+
344
+ type: Literal["workflow_status"] = "workflow_status"
345
+ session_id: str
346
+ status: WorkflowStatus
347
+ execution_mode: ExecutionMode
348
+ details: str | None = None
349
+
350
+ @classmethod
351
+ def make(
352
+ cls,
353
+ session_id: str,
354
+ status: WorkflowStatus,
355
+ mode: ExecutionMode,
356
+ details: str | None = None,
357
+ ) -> "WorkflowStatusNotification":
358
+ """Create a workflow status notification.
359
+
360
+ Parameters
361
+ ----------
362
+ session_id : str
363
+ The session ID.
364
+ status : WorkflowStatus
365
+ The workflow status.
366
+ mode : ExecutionMode
367
+ The execution mode.
368
+ details : str | None
369
+ Additional details about the status.
370
+
371
+ Returns
372
+ -------
373
+ WorkflowStatusNotification
374
+ The created workflow status notification.
375
+ """
376
+ return cls(
377
+ session_id=session_id,
378
+ status=status,
379
+ execution_mode=mode,
380
+ details=details,
381
+ )
382
+
383
+
384
+ class WorkflowOutputNotification(BaseNotification):
385
+ """Notification of workflow output."""
386
+
387
+ type: Literal["workflow_output"] = "workflow_output"
388
+ session_id: str
389
+ stream: Literal["stdout", "stderr"]
390
+ content: str
391
+ output_type: Literal["text", "structured", "debug"] = "text"
392
+
393
+ @classmethod
394
+ def stdout(cls, session_id: str, text: str) -> "WorkflowOutputNotification":
395
+ """Create a workflow output notification for stdout.
396
+
397
+ Parameters
398
+ ----------
399
+ session_id : str
400
+ The session ID.
401
+ text : str
402
+ The output text.
403
+
404
+ Returns
405
+ -------
406
+ WorkflowOutputNotification
407
+ The created workflow output notification.
408
+ """
409
+ return cls(session_id=session_id, stream="stdout", content=text)
410
+
411
+ @classmethod
412
+ def stderr(cls, session_id: str, text: str) -> "WorkflowOutputNotification":
413
+ """Create a workflow output notification for stderr.
414
+
415
+ Parameters
416
+ ----------
417
+ session_id : str
418
+ The session ID.
419
+ text : str
420
+ The output text.
421
+
422
+ Returns
423
+ -------
424
+ WorkflowOutputNotification
425
+ The created workflow output notification.
426
+ """
427
+ return cls(session_id=session_id, stream="stderr", content=text)
428
+
429
+
430
+ class WorkflowEventNotification(BaseNotification):
431
+ """Notification of workflow event (for step-by-step)."""
432
+
433
+ type: Literal["workflow_event"] = "workflow_event"
434
+ session_id: str
435
+ event_data: dict[str, Any]
436
+ event_count: int
437
+ should_break: bool = False
438
+
439
+
440
+ class UserInputRequestNotification(BaseNotification):
441
+ """Notification requesting user input."""
442
+
443
+ type: Literal["input_request"] = "input_request"
444
+ session_id: str
445
+ request_id: str
446
+ prompt: str
447
+ password: bool = False
448
+ timeout: float = 120.0
449
+
450
+
451
+ class BreakpointNotification(BaseNotification):
452
+ """Notification about breakpoint changes."""
453
+
454
+ type: Literal["breakpoint_notification"] = "breakpoint_notification"
455
+ session_id: str
456
+ action: Literal["added", "removed", "cleared", "list"]
457
+ event_type: str | None = None
458
+ breakpoints: list[str] = Field(default_factory=list)
459
+ message: str | None = None
460
+
461
+
462
+ class WorkflowCompletionNotification(BaseNotification):
463
+ """Notification of workflow completion."""
464
+
465
+ type: Literal["workflow_completion"] = "workflow_completion"
466
+ session_id: str
467
+ success: bool
468
+ exit_code: int | None = None
469
+ results: list[dict[str, Any]] = Field(default_factory=list)
470
+ execution_time: float | None = None
471
+ error: str | None = None
472
+
473
+
474
+ class StepDebugNotification(BaseNotification):
475
+ """Notification for step-by-step debug information."""
476
+
477
+ type: Literal["step_debug"] = "step_debug"
478
+ session_id: str
479
+ debug_type: Literal["stats", "help", "error", "info"]
480
+ data: dict[str, Any]
481
+
482
+
483
+ class ConnectionNotification(BaseNotification):
484
+ """Notification about connection status."""
485
+
486
+ type: Literal["connection"] = "connection"
487
+ status: Literal["connected", "disconnected", "error"]
488
+ client_id: str
489
+ server_time: float = Field(default_factory=time.time)
490
+ message: str | None = None
491
+
492
+
493
+ # ========================================
494
+ # SUBPROCESS-SPECIFIC MODELS
495
+ # ========================================
496
+
497
+
498
+ class SubprocessOutputNotification(BaseNotification):
499
+ """Notification from subprocess execution."""
500
+
501
+ type: Literal["subprocess_output"] = "subprocess_output"
502
+ session_id: str
503
+ stream: Literal["stdout", "stderr"]
504
+ content: str
505
+ subprocess_type: Literal["output", "error", "debug"] = "output"
506
+ context: dict[str, Any] = Field(default_factory=dict)
507
+
508
+
509
+ class SubprocessInputRequestNotification(BaseNotification):
510
+ """Input request from subprocess."""
511
+
512
+ type: Literal["subprocess_input_request"] = "subprocess_input_request"
513
+ session_id: str
514
+ request_id: str
515
+ prompt: str
516
+ timeout: float = 120.0
517
+ password: bool = False
518
+ context: dict[str, Any] = Field(default_factory=dict)
519
+
520
+
521
+ class SubprocessCompletionNotification(BaseNotification):
522
+ """Subprocess execution completion."""
523
+
524
+ type: Literal["subprocess_completion"] = "subprocess_completion"
525
+ session_id: str
526
+ success: bool
527
+ exit_code: int
528
+ message: str
529
+ context: dict[str, Any] = Field(default_factory=dict)
530
+
531
+
532
+ # ========================================
533
+ # UNION TYPES FOR MESSAGE PARSING
534
+ # ========================================
535
+
536
+ # Client-to-Server Messages
537
+ ClientMessage = Annotated[
538
+ Union[
539
+ SaveFlowRequest,
540
+ RunWorkflowRequest,
541
+ StepRunWorkflowRequest,
542
+ StepControlRequest,
543
+ BreakpointRequest,
544
+ UserInputResponse,
545
+ StopWorkflowRequest,
546
+ ConvertWorkflowRequest,
547
+ UploadFileRequest,
548
+ PingRequest,
549
+ GetStatusRequest,
550
+ ],
551
+ Field(discriminator="type"),
552
+ ]
553
+
554
+ # Server-to-Client Messages
555
+ ServerMessage = Annotated[
556
+ Union[
557
+ # Responses
558
+ SaveFlowResponse,
559
+ RunWorkflowResponse,
560
+ StepRunWorkflowResponse,
561
+ StopWorkflowResponse,
562
+ StepControlResponse,
563
+ BreakpointResponse,
564
+ BreakpointNotification,
565
+ ConvertWorkflowResponse,
566
+ UploadFileResponse,
567
+ PongResponse,
568
+ StatusResponse,
569
+ ErrorResponse,
570
+ # Notifications
571
+ WorkflowStatusNotification,
572
+ WorkflowOutputNotification,
573
+ WorkflowEventNotification,
574
+ UserInputRequestNotification,
575
+ WorkflowCompletionNotification,
576
+ StepDebugNotification,
577
+ ConnectionNotification,
578
+ SubprocessOutputNotification,
579
+ SubprocessInputRequestNotification,
580
+ SubprocessCompletionNotification,
581
+ ],
582
+ Field(discriminator="type"),
583
+ ]
584
+
585
+
586
+ class _ServerMessageWrapper(BaseModel):
587
+ """Wrapper for server messages to handle discriminators."""
588
+
589
+ # noinspection PyTypeHints
590
+ message: ServerMessage
591
+
592
+
593
+ class _ClientMessageWrapper(BaseModel):
594
+ """Wrapper for client messages to handle discriminators."""
595
+
596
+ # noinspection PyTypeHints
597
+ message: ClientMessage
598
+
599
+
600
+ # All messages
601
+ WaldiezWsMessage = Annotated[
602
+ Union[ClientMessage, ServerMessage],
603
+ Field(discriminator="type"),
604
+ ]
605
+
606
+
607
+ # ========================================
608
+ # UTILITY FUNCTIONS
609
+ # ========================================
610
+
611
+
612
+ # noinspection PyTypeHints,DuplicatedCode
613
+ def parse_client_message(data: str | dict[str, Any]) -> ClientMessage:
614
+ """Parse client message from JSON string or dict.
615
+
616
+ Parameters
617
+ ----------
618
+ data : str | dict[str, Any]
619
+ Message data
620
+
621
+ Returns
622
+ -------
623
+ ClientMessage
624
+ Parsed client message
625
+
626
+ Raises
627
+ ------
628
+ ValueError
629
+ If message cannot be parsed
630
+ """
631
+ if isinstance(data, str):
632
+ try:
633
+ data = json.loads(data)
634
+ except json.JSONDecodeError as e:
635
+ raise ValueError(f"Invalid message format: {e}") from e
636
+ try:
637
+ return _ClientMessageWrapper.model_validate({"message": data}).message
638
+ except ValidationError as e:
639
+ raise ValueError(f"Invalid client message: {e}") from e
640
+
641
+
642
+ # noinspection PyTypeHints,DuplicatedCode
643
+ def parse_server_message(data: str | dict[str, Any]) -> ServerMessage:
644
+ """Parse server message from JSON string or dict.
645
+
646
+ Parameters
647
+ ----------
648
+ data : str | dict[str, Any]
649
+ Message data
650
+
651
+ Returns
652
+ -------
653
+ ServerMessage
654
+ Parsed server message
655
+
656
+ Raises
657
+ ------
658
+ ValueError
659
+ If message cannot be parsed
660
+ """
661
+ if isinstance(data, str):
662
+ try:
663
+ data = json.loads(data)
664
+ except json.JSONDecodeError as e:
665
+ raise ValueError(f"Invalid message format: {e}") from e
666
+ try:
667
+ return _ServerMessageWrapper.model_validate({"message": data}).message
668
+ except ValidationError as e:
669
+ raise ValueError(f"Invalid server message: {e}") from e
670
+
671
+
672
+ def create_error_response(
673
+ error_message: str,
674
+ error_code: int = 500,
675
+ error_type: str = "InternalError",
676
+ request_id: str | None = None,
677
+ details: dict[str, Any] | None = None,
678
+ ) -> ErrorResponse:
679
+ """Create standardized error response.
680
+
681
+ Parameters
682
+ ----------
683
+ error_message : str
684
+ Error message
685
+ error_code : int
686
+ Error code
687
+ error_type : str
688
+ Error type
689
+ request_id : str | None
690
+ Request ID if this is response to a request
691
+ details : dict[str, Any] | None
692
+ Additional error details
693
+
694
+ Returns
695
+ -------
696
+ ErrorResponse
697
+ Standardized error response
698
+ """
699
+ return ErrorResponse(
700
+ request_id=request_id,
701
+ error_code=error_code,
702
+ error_type=error_type,
703
+ error=error_message,
704
+ details=details or {},
705
+ )
706
+
707
+
708
+ def create_session_id() -> str:
709
+ """Create a new session ID.
710
+
711
+ Returns
712
+ -------
713
+ str
714
+ New session ID
715
+ """
716
+ return f"session_{uuid.uuid4().hex[:12]}"
717
+
718
+
719
+ # ========================================
720
+ # MESSAGE TYPE CONSTANTS
721
+ # ========================================
722
+
723
+ CLIENT_MESSAGE_TYPES = {
724
+ "save_flow",
725
+ "run_workflow",
726
+ "step_run_workflow",
727
+ "step_control",
728
+ "breakpoint_control",
729
+ "user_input",
730
+ "stop_workflow",
731
+ "convert_workflow",
732
+ "upload_file",
733
+ "ping",
734
+ "get_status",
735
+ }
736
+
737
+ SERVER_MESSAGE_TYPES = {
738
+ # Responses
739
+ "save_flow_response",
740
+ "run_workflow_response",
741
+ "step_run_workflow_response",
742
+ "step_control_response",
743
+ "breakpoint_response",
744
+ "convert_workflow_response",
745
+ "upload_file_response",
746
+ "pong",
747
+ "status_response",
748
+ "error",
749
+ # Notifications
750
+ "workflow_status",
751
+ "workflow_output",
752
+ "workflow_event",
753
+ "input_request",
754
+ "breakpoint_notification",
755
+ "workflow_completion",
756
+ "step_debug",
757
+ "connection",
758
+ "subprocess_output",
759
+ "subprocess_input_request",
760
+ "subprocess_completion",
761
+ }
762
+
763
+ ALL_MESSAGE_TYPES = CLIENT_MESSAGE_TYPES | SERVER_MESSAGE_TYPES
764
+
765
+
766
+ class SessionState(BaseModel):
767
+ """State information for a workflow session."""
768
+
769
+ session_id: str
770
+ client_id: str
771
+ status: WorkflowStatus
772
+ execution_mode: ExecutionMode
773
+ start_time: int = Field(default_factory=time.monotonic_ns)
774
+ end_time: int | None = None
775
+ metadata: dict[str, Any] = Field(default_factory=dict)
776
+
777
+ # Runtime fields (not serialized)
778
+ runner: Any = Field(default=None, exclude=True)
779
+ temp_file: Path | None = Field(default=None, exclude=True)
780
+
781
+ @property
782
+ def duration(self) -> float:
783
+ """Get session duration in seconds."""
784
+ end = self.end_time or time.monotonic_ns()
785
+ return (end - self.start_time) / 1_000_000_000
786
+
787
+ @property
788
+ def is_active(self) -> bool:
789
+ """Check if session is currently active."""
790
+ return self.status in {
791
+ WorkflowStatus.STARTING,
792
+ WorkflowStatus.RUNNING,
793
+ WorkflowStatus.PAUSED,
794
+ WorkflowStatus.STEP_WAITING,
795
+ WorkflowStatus.INPUT_WAITING,
796
+ }
797
+
798
+ @property
799
+ def is_completed(self) -> bool:
800
+ """Check if session has completed (successfully or not)."""
801
+ return self.status in {
802
+ WorkflowStatus.COMPLETED,
803
+ WorkflowStatus.FAILED,
804
+ WorkflowStatus.CANCELLED,
805
+ }
806
+
807
+ def update_status(self, new_status: WorkflowStatus) -> None:
808
+ """Update session status and set end time if completed.
809
+
810
+ Parameters
811
+ ----------
812
+ new_status : WorkflowStatus
813
+ The new status to set.
814
+ """
815
+ self.status = new_status
816
+ if self.is_completed and not self.end_time:
817
+ self.end_time = time.monotonic_ns()
818
+
819
+ def get_execution_summary(self) -> dict[str, Any]:
820
+ """Get a summary of session execution.
821
+
822
+ Returns
823
+ -------
824
+ dict[str, Any]
825
+ A dictionary containing session execution summary.
826
+ """
827
+ return {
828
+ "session_id": self.session_id,
829
+ "client_id": self.client_id,
830
+ "status": self.status.value,
831
+ "execution_mode": self.execution_mode.value,
832
+ "duration_seconds": self.duration,
833
+ "start_time": self.start_time,
834
+ "end_time": self.end_time,
835
+ "is_active": self.is_active,
836
+ "is_completed": self.is_completed,
837
+ }
838
+
839
+
840
+ # noinspection TryExceptPass,PyBroadException
841
+ class WorkflowSession:
842
+ """Enhanced session wrapper with runtime management capabilities."""
843
+
844
+ def __init__(
845
+ self,
846
+ session_state: SessionState,
847
+ runner: Any = None,
848
+ temp_file: Path | None = None,
849
+ ):
850
+ """Initialize workflow session.
851
+
852
+ Parameters
853
+ ----------
854
+ session_state : SessionState
855
+ The session state data
856
+ runner : Any, optional
857
+ The workflow runner instance
858
+ temp_file : Path, optional
859
+ Temporary file path for cleanup
860
+ """
861
+ self._state = session_state
862
+ self.runner = runner
863
+ self._temp_file = temp_file
864
+ self._created_at = time.monotonic_ns()
865
+ self._last_accessed = time.monotonic_ns()
866
+ self._access_count = 0
867
+
868
+ @property
869
+ def start_time(self) -> int:
870
+ """Get the start time of the session."""
871
+ return self._state.start_time
872
+
873
+ @property
874
+ def state(self) -> SessionState:
875
+ """Get the session state."""
876
+ self._last_accessed = time.monotonic_ns()
877
+ self._access_count += 1
878
+ return self._state
879
+
880
+ @property
881
+ def raw_state(self) -> SessionState:
882
+ """Get the raw session state."""
883
+ return self._state
884
+
885
+ @property
886
+ def session_id(self) -> str:
887
+ """Get session ID."""
888
+ return self._state.session_id
889
+
890
+ @property
891
+ def client_id(self) -> str:
892
+ """Get client ID."""
893
+ return self._state.client_id
894
+
895
+ @property
896
+ def status(self) -> WorkflowStatus:
897
+ """Get current status."""
898
+ return self._state.status
899
+
900
+ @property
901
+ def execution_mode(self) -> ExecutionMode:
902
+ """Get execution mode."""
903
+ return self._state.execution_mode
904
+
905
+ @property
906
+ def temp_file(self) -> Path | None:
907
+ """Get temporary file path."""
908
+ return self._temp_file
909
+
910
+ @property
911
+ def metadata(self) -> dict[str, Any]:
912
+ """Get session metadata."""
913
+ return self._state.metadata
914
+
915
+ @property
916
+ def last_accessed(self) -> float:
917
+ """Get last access time."""
918
+ return self._last_accessed
919
+
920
+ @property
921
+ def access_count(self) -> int:
922
+ """Get access count."""
923
+ return self._access_count
924
+
925
+ def update_status(self, new_status: WorkflowStatus) -> None:
926
+ """Update session status.
927
+
928
+ Parameters
929
+ ----------
930
+ new_status : WorkflowStatus
931
+ The new status to set.
932
+ """
933
+ self._state.update_status(new_status)
934
+ self._access_count += 1
935
+ self._last_accessed = time.monotonic_ns()
936
+
937
+ def cleanup(self) -> None:
938
+ """Cleanup session resources."""
939
+ # Stop runner if still running
940
+ if self.runner and hasattr(self.runner, "stop"):
941
+ try:
942
+ self.runner.stop()
943
+ except Exception:
944
+ pass # Best effort cleanup
945
+
946
+ # Remove temporary file
947
+ if self._temp_file and self._temp_file.exists():
948
+ try:
949
+ self._temp_file.unlink()
950
+ except Exception:
951
+ pass # Best effort cleanup
952
+
953
+ def to_dict(self) -> dict[str, Any]:
954
+ """Convert session to dictionary representation.
955
+
956
+ Returns
957
+ -------
958
+ dict[str, Any]
959
+ The dictionary representation of the session.
960
+ """
961
+ base_dict = self._state.get_execution_summary()
962
+ base_dict.update(
963
+ {
964
+ "created_at": self._created_at,
965
+ "last_accessed": self._last_accessed,
966
+ "access_count": self._access_count,
967
+ "has_runner": self.runner is not None,
968
+ "has_temp_file": self._temp_file is not None,
969
+ }
970
+ )
971
+ return base_dict