toolplane-python-client 0.1.0__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.
Files changed (81) hide show
  1. toolplane/__init__.py +106 -0
  2. toolplane/common/__init__.py +93 -0
  3. toolplane/common/base_config.py +129 -0
  4. toolplane/common/base_connection_manager.py +171 -0
  5. toolplane/common/base_session_manager.py +321 -0
  6. toolplane/common/base_tool_manager.py +347 -0
  7. toolplane/common/constants.py +47 -0
  8. toolplane/common/utils.py +310 -0
  9. toolplane/core/__init__.py +67 -0
  10. toolplane/core/config.py +107 -0
  11. toolplane/core/connection.py +285 -0
  12. toolplane/core/errors.py +298 -0
  13. toolplane/core/machine.py +480 -0
  14. toolplane/core/request.py +775 -0
  15. toolplane/core/session.py +332 -0
  16. toolplane/core/session_context.py +514 -0
  17. toolplane/core/task.py +130 -0
  18. toolplane/core/tool.py +329 -0
  19. toolplane/http_core/__init__.py +37 -0
  20. toolplane/http_core/http_config.py +97 -0
  21. toolplane/http_core/http_connection.py +409 -0
  22. toolplane/http_core/http_machine.py +298 -0
  23. toolplane/http_core/http_request.py +748 -0
  24. toolplane/http_core/http_session.py +348 -0
  25. toolplane/http_core/http_session_context.py +491 -0
  26. toolplane/http_core/http_task.py +101 -0
  27. toolplane/http_core/http_tool.py +400 -0
  28. toolplane/interfaces/__init__.py +27 -0
  29. toolplane/interfaces/client_interface.py +122 -0
  30. toolplane/interfaces/connection_interface.py +193 -0
  31. toolplane/interfaces/event_interface.py +290 -0
  32. toolplane/interfaces/request_interface.py +439 -0
  33. toolplane/interfaces/session_interface.py +288 -0
  34. toolplane/interfaces/tool_interface.py +441 -0
  35. toolplane/proto/__init__.py +0 -0
  36. toolplane/proto/service_pb2.py +315 -0
  37. toolplane/proto/service_pb2_grpc.py +2240 -0
  38. toolplane/provider_cli.py +268 -0
  39. toolplane/provider_registry.py +77 -0
  40. toolplane/provider_runtime.py +302 -0
  41. toolplane/toolkits/__init__.py +0 -0
  42. toolplane/toolkits/standalone_tools/__init__.py +0 -0
  43. toolplane/toolkits/standalone_tools/create_directory.py +94 -0
  44. toolplane/toolkits/standalone_tools/create_file.py +124 -0
  45. toolplane/toolkits/standalone_tools/file_search.py +229 -0
  46. toolplane/toolkits/standalone_tools/grep_search.py +372 -0
  47. toolplane/toolkits/standalone_tools/launcher.py +146 -0
  48. toolplane/toolkits/standalone_tools/list_dir.py +395 -0
  49. toolplane/toolkits/standalone_tools/read_file.py +346 -0
  50. toolplane/toolkits/standalone_tools/replace_string_in_file.py +407 -0
  51. toolplane/toolkits/standalone_tools/run_tests.py +66 -0
  52. toolplane/toolkits/standalone_tools/semantic_search.py +485 -0
  53. toolplane/toolkits/standalone_tools/standalone_toolkit.py +979 -0
  54. toolplane/toolkits/standalone_tools/test_failure_analysis.py +618 -0
  55. toolplane/toolkits/standalone_tools/test_standalone_toolkit.py +517 -0
  56. toolplane/toolkits/swe/__init__.py +35 -0
  57. toolplane/toolkits/swe/create_directory.py +15 -0
  58. toolplane/toolkits/swe/create_file.py +15 -0
  59. toolplane/toolkits/swe/descriptions.py +273 -0
  60. toolplane/toolkits/swe/execute_bash.py +93 -0
  61. toolplane/toolkits/swe/file_editor.py +775 -0
  62. toolplane/toolkits/swe/file_search.py +16 -0
  63. toolplane/toolkits/swe/finish.py +50 -0
  64. toolplane/toolkits/swe/grep_search.py +19 -0
  65. toolplane/toolkits/swe/list_dir.py +407 -0
  66. toolplane/toolkits/swe/read_file.py +18 -0
  67. toolplane/toolkits/swe/replace_string_in_file.py +17 -0
  68. toolplane/toolkits/swe/search.py +260 -0
  69. toolplane/toolkits/swe/semantic_search.py +20 -0
  70. toolplane/toolkits/swe/str_replace_editor.py +647 -0
  71. toolplane/toolkits/swe/submit.py +29 -0
  72. toolplane/toolkits/swe/swe_toolkit.py +1296 -0
  73. toolplane/toolplane_client.py +686 -0
  74. toolplane/toolplane_http_client.py +681 -0
  75. toolplane/utils/__init__.py +3 -0
  76. toolplane/utils/schema.py +146 -0
  77. toolplane_python_client-0.1.0.dist-info/METADATA +543 -0
  78. toolplane_python_client-0.1.0.dist-info/RECORD +81 -0
  79. toolplane_python_client-0.1.0.dist-info/WHEEL +5 -0
  80. toolplane_python_client-0.1.0.dist-info/entry_points.txt +2 -0
  81. toolplane_python_client-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,681 @@
1
+ """Modular HTTP Toolplane client implementation."""
2
+
3
+ import logging
4
+ import threading
5
+ import time
6
+ from typing import Any, Callable, Dict, List, Optional
7
+
8
+ from .common import (
9
+ generate_session_id,
10
+ validate_session_name,
11
+ )
12
+ from .http_core import (
13
+ ConnectionError,
14
+ HTTPClientConfig,
15
+ HTTPConnectionManager,
16
+ HTTPMachineManager,
17
+ HTTPRequestManager,
18
+ HTTPSessionContext,
19
+ HTTPSessionManager,
20
+ HTTPTaskManager,
21
+ HTTPToolManager,
22
+ ToolplaneError,
23
+ )
24
+
25
+ logger = logging.getLogger(__name__)
26
+
27
+
28
+ class ToolplaneHTTP:
29
+ """
30
+ Modular HTTP Toolplane client for registering and executing tools.
31
+ Supports multi-session mode with explicit session management.
32
+ """
33
+
34
+ def __init__(
35
+ self,
36
+ server_host: str = "localhost",
37
+ server_port: int = 8080,
38
+ session_ids: Optional[List[str]] = None,
39
+ api_key: Optional[str] = None,
40
+ user_id: Optional[str] = None,
41
+ session_name: Optional[str] = None,
42
+ session_description: Optional[str] = None,
43
+ session_namespace: Optional[str] = None,
44
+ max_buffer_size: int = 4 * 1024 * 1024, # 4MB default
45
+ max_retries: int = 3,
46
+ request_timeout: int = 30, # 30 seconds
47
+ retry_backoff_ms: int = 250, # Start with 250ms backoff
48
+ heartbeat_interval: int = 60,
49
+ max_workers: int = 10,
50
+ poll_interval: float = 1.0,
51
+ event_emitter: Optional[Any] = None,
52
+ ):
53
+ """Initialize HTTP Toolplane client."""
54
+ # Create configuration
55
+ self.config = HTTPClientConfig(
56
+ server_host=server_host,
57
+ server_port=server_port,
58
+ api_key=api_key,
59
+ user_id=user_id,
60
+ session_name=session_name,
61
+ session_description=session_description,
62
+ session_namespace=session_namespace,
63
+ max_buffer_size=max_buffer_size,
64
+ max_retries=max_retries,
65
+ request_timeout=request_timeout,
66
+ retry_backoff_ms=retry_backoff_ms,
67
+ heartbeat_interval=heartbeat_interval,
68
+ max_workers=max_workers,
69
+ poll_interval=poll_interval,
70
+ )
71
+
72
+ # Initialize managers
73
+ self.connection_manager = HTTPConnectionManager(self.config)
74
+ self.machine_manager = HTTPMachineManager(
75
+ self.connection_manager, event_emitter
76
+ )
77
+ self.tool_manager = HTTPToolManager(self.connection_manager)
78
+ self.request_manager = HTTPRequestManager(self.connection_manager, max_workers)
79
+ self.task_manager = HTTPTaskManager(self.connection_manager)
80
+ self.session_manager = HTTPSessionManager(self.connection_manager)
81
+
82
+ # Client state
83
+ self.running = False
84
+ self.session_ids = (
85
+ session_ids or []
86
+ ) # Changed from single session_id to session_ids list
87
+ self._main_thread: Optional[threading.Thread] = None
88
+ self._provider_runtime = None
89
+
90
+ def connect(self) -> bool:
91
+ """Connect to HTTP server."""
92
+ try:
93
+ success = self.connection_manager.connect()
94
+ if success:
95
+ self._initialize_sessions(register_machine=False)
96
+ return success
97
+ except Exception as e:
98
+ raise ConnectionError(f"Failed to connect: {e}")
99
+
100
+ def disconnect(self):
101
+ """Disconnect from HTTP server."""
102
+ self.stop()
103
+ try:
104
+ self.request_manager.stop_polling()
105
+ except Exception:
106
+ pass
107
+ try:
108
+ self.request_manager.stop_lease_renewal()
109
+ except Exception:
110
+ pass
111
+ try:
112
+ self.request_manager.shutdown()
113
+ except Exception:
114
+ pass
115
+ try:
116
+ self.machine_manager.stop_heartbeat()
117
+ except Exception:
118
+ pass
119
+ self.machine_manager.cleanup_all()
120
+ self.tool_manager.cleanup_all()
121
+ self.session_manager.cleanup_all()
122
+ self.connection_manager.disconnect()
123
+
124
+ def _initialize_sessions(self, register_machine: bool = False):
125
+ """Initialize specified sessions."""
126
+ for session_id in list(self.session_ids):
127
+ try:
128
+ self.ensure_session_context(
129
+ session_id,
130
+ create_if_missing=True,
131
+ register_machine=register_machine,
132
+ )
133
+ except Exception as e:
134
+ logger.warning("Failed to initialize session %s: %s", session_id, e)
135
+
136
+ def ensure_session_context(
137
+ self,
138
+ session_id: str,
139
+ create_if_missing: bool = False,
140
+ register_machine: bool = False,
141
+ ) -> HTTPSessionContext:
142
+ """Ensure a local session context exists without implying provider startup."""
143
+ if not session_id or not isinstance(session_id, str):
144
+ raise ToolplaneError(f"Invalid session ID: {session_id}")
145
+
146
+ existing_context = self.session_manager.get_session_context(session_id)
147
+ if existing_context:
148
+ if register_machine and not getattr(existing_context, "machine_id", None):
149
+ if not existing_context.register_machine():
150
+ raise ToolplaneError(
151
+ f"Failed to register machine for session {session_id}"
152
+ )
153
+ if session_id not in self.session_ids:
154
+ self.session_ids.append(session_id)
155
+ return existing_context
156
+
157
+ resolved_session_id = session_id
158
+ if not self.session_manager.get_session(session_id):
159
+ if not create_if_missing:
160
+ raise ToolplaneError(f"Session {session_id} does not exist")
161
+ resolved_session_id = self.session_manager.create_session(
162
+ session_id=session_id,
163
+ user_id=self.config.user_id,
164
+ name=self.config.session_name,
165
+ description=self.config.session_description,
166
+ namespace=self.config.session_namespace,
167
+ )
168
+
169
+ context = HTTPSessionContext(
170
+ resolved_session_id,
171
+ self.connection_manager,
172
+ self.machine_manager,
173
+ self.tool_manager,
174
+ self.request_manager,
175
+ self.session_manager,
176
+ )
177
+ self.session_manager.register_session_context(resolved_session_id, context)
178
+ if resolved_session_id not in self.session_ids:
179
+ self.session_ids.append(resolved_session_id)
180
+
181
+ if register_machine and not context.register_machine():
182
+ raise ToolplaneError(
183
+ f"Failed to register machine for session {resolved_session_id}"
184
+ )
185
+
186
+ return context
187
+
188
+ def create_session(
189
+ self,
190
+ session_id: Optional[str] = None,
191
+ user_id: Optional[str] = None,
192
+ name: Optional[str] = None,
193
+ description: Optional[str] = None,
194
+ namespace: Optional[str] = None,
195
+ register_machine: bool = False,
196
+ ) -> HTTPSessionContext:
197
+ """Create a new session."""
198
+ if not self.connection_manager.connected:
199
+ if not self.connect():
200
+ raise ConnectionError("Failed to connect to server")
201
+
202
+ # Validate session name if provided
203
+ if name and not validate_session_name(name):
204
+ raise ToolplaneError(f"Invalid session name: {name}")
205
+
206
+ # Generate session ID if not provided
207
+ if not session_id:
208
+ session_id = generate_session_id()
209
+
210
+ try:
211
+ # Create session on server
212
+ created_session_id = self.session_manager.create_session(
213
+ session_id=session_id,
214
+ user_id=user_id or self.config.user_id,
215
+ name=name or self.config.session_name,
216
+ description=description or self.config.session_description,
217
+ namespace=namespace or self.config.session_namespace,
218
+ )
219
+
220
+ logger.info("Created new session: %s", created_session_id)
221
+
222
+ context = self.ensure_session_context(
223
+ created_session_id,
224
+ create_if_missing=False,
225
+ register_machine=register_machine,
226
+ )
227
+ return context
228
+
229
+ except Exception as e:
230
+ raise ToolplaneError(f"Failed to create session: {e}")
231
+
232
+ def get_session(self, session_id: str) -> Optional[HTTPSessionContext]:
233
+ """Get session context by ID."""
234
+ return self.session_manager.get_session_context(session_id)
235
+
236
+ def list_sessions(self) -> List[HTTPSessionContext]:
237
+ """List all session contexts."""
238
+ return self.session_manager.list_session_contexts()
239
+
240
+ def get_primary_session_context(self) -> Optional[HTTPSessionContext]:
241
+ """Get primary session context."""
242
+ if self.session_ids:
243
+ return self.get_session(self.session_ids[0])
244
+ return None
245
+
246
+ def tool(
247
+ self,
248
+ session_id: str,
249
+ name: Optional[str] = None,
250
+ description: Optional[str] = None,
251
+ stream: bool = False,
252
+ tags: Optional[List[str]] = None,
253
+ ):
254
+ """Decorator to register a tool for a session."""
255
+
256
+ def decorator(func):
257
+ context = self.get_session(session_id)
258
+ if not context:
259
+ raise ToolplaneError(f"Session {session_id} not found")
260
+
261
+ context.register_tool(
262
+ name or func.__name__,
263
+ func,
264
+ description=description,
265
+ stream=stream,
266
+ tags=tags or [],
267
+ )
268
+ return func
269
+
270
+ return decorator
271
+
272
+ def invoke(
273
+ self,
274
+ tool_name: str,
275
+ session_id: str,
276
+ timeout_seconds: int = 0,
277
+ wait_timeout: Optional[int] = None,
278
+ **params,
279
+ ) -> Any:
280
+ """Invoke a tool in a session and return the tool result value.
281
+
282
+ timeout_seconds sets the request's absolute per-attempt execution
283
+ timeout on the wire (0 keeps the server default); wait_timeout bounds
284
+ the local wait (default: timeout_seconds + 15, else 60).
285
+ """
286
+ context = self.get_session(session_id)
287
+ if not context:
288
+ raise ToolplaneError(f"Session {session_id} not found")
289
+
290
+ return context.invoke(
291
+ tool_name,
292
+ timeout_seconds=timeout_seconds,
293
+ wait_timeout=wait_timeout,
294
+ **params,
295
+ )
296
+
297
+ async def ainvoke(
298
+ self, tool_name: str, session_id: str, timeout_seconds: int = 0, **params
299
+ ) -> str:
300
+ """Submit a tool invocation without blocking; awaits the request ID."""
301
+ context = self.get_session(session_id)
302
+ if not context:
303
+ raise ToolplaneError(f"Session {session_id} not found")
304
+
305
+ return await context.ainvoke(
306
+ tool_name, timeout_seconds=timeout_seconds, **params
307
+ )
308
+
309
+ def stream(
310
+ self,
311
+ tool_name: str,
312
+ callback: Callable[[Any, bool], None],
313
+ session_id: str,
314
+ **params,
315
+ ) -> List[Any]:
316
+ """Stream tool execution."""
317
+ context = self.get_session(session_id)
318
+ if not context:
319
+ raise ToolplaneError(f"Session {session_id} not found")
320
+
321
+ return context.stream(tool_name, callback, **params)
322
+
323
+ async def astream(
324
+ self,
325
+ tool_name: str,
326
+ callback: Callable[[Any, bool], None],
327
+ session_id: str,
328
+ **params,
329
+ ) -> List[Any]:
330
+ """Awaitable stream: resolves with the collected chunks."""
331
+ context = self.get_session(session_id)
332
+ if not context:
333
+ raise ToolplaneError(f"Session {session_id} not found")
334
+
335
+ return await context.astream(tool_name, callback, **params)
336
+
337
+ def get_available_tools(self, session_id: str) -> Dict[str, Any]:
338
+ """Get available tools for a session."""
339
+ context = self.get_session(session_id)
340
+ if not context:
341
+ raise ToolplaneError(f"Session {session_id} not found")
342
+
343
+ return context.get_available_tools()
344
+
345
+ def list_tools(self, session_id: str) -> List[Dict[str, Any]]:
346
+ """List tools for a session."""
347
+ if not self.connection_manager.connected:
348
+ if not self.connect():
349
+ raise ConnectionError("Failed to connect to server")
350
+ return self.tool_manager.list_tools(session_id)
351
+
352
+ def get_tool_by_id(self, session_id: str, tool_id: str) -> Dict[str, Any]:
353
+ """Get a tool by ID."""
354
+ if not self.connection_manager.connected:
355
+ if not self.connect():
356
+ raise ConnectionError("Failed to connect to server")
357
+ return self.tool_manager.get_tool_by_id(session_id, tool_id)
358
+
359
+ def get_tool_by_name(self, session_id: str, tool_name: str) -> Dict[str, Any]:
360
+ """Get a tool by name."""
361
+ if not self.connection_manager.connected:
362
+ if not self.connect():
363
+ raise ConnectionError("Failed to connect to server")
364
+ return self.tool_manager.get_tool_by_name(session_id, tool_name)
365
+
366
+ def delete_tool(self, session_id: str, tool_id: str) -> bool:
367
+ """Delete a tool by ID."""
368
+ if not self.connection_manager.connected:
369
+ if not self.connect():
370
+ raise ConnectionError("Failed to connect to server")
371
+ return self.tool_manager.delete_tool(session_id, tool_id)
372
+
373
+ def get_request_status(self, session_id: str, request_id: str) -> Dict[str, Any]:
374
+ """Get request status (session-scoped, like every other facade method)."""
375
+ context = self.get_session(session_id)
376
+ if not context:
377
+ raise ToolplaneError(f"Session {session_id} not found")
378
+
379
+ return context.get_request_status(request_id)
380
+
381
+ # Session admin helpers (admin scope — Python-only; not portable across SDKs)
382
+ def list_user_sessions(
383
+ self, user_id: str, page_size: int = 10, page_token: str = "", filter: str = ""
384
+ ) -> Dict[str, Any]:
385
+ """List user sessions with pagination and filtering.
386
+
387
+ page_token is the opaque cursor returned by the previous page.
388
+ """
389
+ if not self.connection_manager.connected:
390
+ if not self.connect():
391
+ raise ConnectionError("Failed to connect to server")
392
+ return self.session_manager.list_user_sessions(
393
+ user_id, page_size, page_token, filter
394
+ )
395
+
396
+ def bulk_delete_sessions(
397
+ self, user_id: str, session_ids: List[str] = None, filter: str = ""
398
+ ) -> Dict[str, Any]:
399
+ """Bulk delete sessions for a user."""
400
+ if not self.connection_manager.connected:
401
+ if not self.connect():
402
+ raise ConnectionError("Failed to connect to server")
403
+ return self.session_manager.bulk_delete_sessions(user_id, session_ids, filter)
404
+
405
+ def get_session_stats(self, user_id: str) -> Dict[str, int]:
406
+ """Get session statistics for a user."""
407
+ if not self.connection_manager.connected:
408
+ if not self.connect():
409
+ raise ConnectionError("Failed to connect to server")
410
+ return self.session_manager.get_session_stats(user_id)
411
+
412
+ def invalidate_session(self, session_id: str, reason: str = "") -> bool:
413
+ """Invalidate a session."""
414
+ if not self.connection_manager.connected:
415
+ if not self.connect():
416
+ raise ConnectionError("Failed to connect to server")
417
+ return self.session_manager.invalidate_session(session_id, reason)
418
+
419
+ def update_session(
420
+ self,
421
+ session_id: str,
422
+ name: Optional[str] = None,
423
+ description: Optional[str] = None,
424
+ namespace: Optional[str] = None,
425
+ ) -> Dict[str, Any]:
426
+ """Update session metadata."""
427
+ if not self.connection_manager.connected:
428
+ if not self.connect():
429
+ raise ConnectionError("Failed to connect to server")
430
+ return self.session_manager.update_session(
431
+ session_id=session_id,
432
+ name=name,
433
+ description=description,
434
+ namespace=namespace,
435
+ )
436
+
437
+ def create_request(
438
+ self,
439
+ session_id: str,
440
+ tool_name: str,
441
+ input_data: str,
442
+ idempotency_key: str = "",
443
+ ) -> str:
444
+ """Create a new request in a session.
445
+
446
+ idempotency_key, when set, dedups creates within the session:
447
+ retrying with the same key returns the original request.
448
+ """
449
+ if not self.connection_manager.connected:
450
+ if not self.connect():
451
+ raise ConnectionError("Failed to connect to server")
452
+ return self.request_manager.create_request(
453
+ session_id, tool_name, input_data, idempotency_key=idempotency_key
454
+ )
455
+
456
+ def list_requests(
457
+ self,
458
+ session_id: str,
459
+ status: str = "",
460
+ tool_name: str = "",
461
+ limit: int = 10,
462
+ page_token: str = "",
463
+ ) -> List[Dict[str, Any]]:
464
+ """List requests in a session.
465
+
466
+ page_token is the opaque cursor returned by the previous page.
467
+ """
468
+ if not self.connection_manager.connected:
469
+ if not self.connect():
470
+ raise ConnectionError("Failed to connect to server")
471
+ return self.request_manager.list_requests(
472
+ session_id=session_id,
473
+ status=status,
474
+ tool_name=tool_name,
475
+ limit=limit,
476
+ page_token=page_token,
477
+ )
478
+
479
+ def list_requests_page(
480
+ self,
481
+ session_id: str,
482
+ status: str = "",
483
+ tool_name: str = "",
484
+ limit: int = 10,
485
+ page_token: str = "",
486
+ ) -> Dict[str, Any]:
487
+ """List one page of requests, with the continuation cursor.
488
+
489
+ Returns {"requests", "next_page_token", "total_size"};
490
+ next_page_token is empty on the last page.
491
+ """
492
+ if not self.connection_manager.connected:
493
+ if not self.connect():
494
+ raise ConnectionError("Failed to connect to server")
495
+ return self.request_manager.list_requests_page(
496
+ session_id=session_id,
497
+ status=status,
498
+ tool_name=tool_name,
499
+ limit=limit,
500
+ page_token=page_token,
501
+ )
502
+
503
+ def cancel_request(self, session_id: str, request_id: str) -> bool:
504
+ """Cancel a request in a session."""
505
+ if not self.connection_manager.connected:
506
+ if not self.connect():
507
+ raise ConnectionError("Failed to connect to server")
508
+ return self.request_manager.cancel_request(session_id, request_id)
509
+
510
+ def create_task(
511
+ self,
512
+ session_id: str,
513
+ tool_name: str,
514
+ input_data: str,
515
+ ) -> Dict[str, Any]:
516
+ """Create a new task in a session."""
517
+ if not self.connection_manager.connected:
518
+ if not self.connect():
519
+ raise ConnectionError("Failed to connect to server")
520
+ return self.task_manager.create_task(session_id, tool_name, input_data)
521
+
522
+ def get_task(self, session_id: str, task_id: str) -> Dict[str, Any]:
523
+ """Get a task by ID in a session."""
524
+ if not self.connection_manager.connected:
525
+ if not self.connect():
526
+ raise ConnectionError("Failed to connect to server")
527
+ return self.task_manager.get_task(session_id, task_id)
528
+
529
+ def list_tasks(self, session_id: str) -> List[Dict[str, Any]]:
530
+ """List tasks in a session."""
531
+ if not self.connection_manager.connected:
532
+ if not self.connect():
533
+ raise ConnectionError("Failed to connect to server")
534
+ return self.task_manager.list_tasks(session_id)
535
+
536
+ def cancel_task(self, session_id: str, task_id: str) -> bool:
537
+ """Cancel a task in a session."""
538
+ if not self.connection_manager.connected:
539
+ if not self.connect():
540
+ raise ConnectionError("Failed to connect to server")
541
+ return self.task_manager.cancel_task(session_id, task_id)
542
+
543
+ def list_machines(self, session_id: str) -> List[Dict[str, Any]]:
544
+ """List machines in a session."""
545
+ if not self.connection_manager.connected:
546
+ if not self.connect():
547
+ raise ConnectionError("Failed to connect to server")
548
+ return self.machine_manager.list_machines(session_id)
549
+
550
+ def get_machine(self, session_id: str, machine_id: str) -> Dict[str, Any]:
551
+ """Get a machine by ID."""
552
+ if not self.connection_manager.connected:
553
+ if not self.connect():
554
+ raise ConnectionError("Failed to connect to server")
555
+ return self.machine_manager.get_machine(session_id, machine_id)
556
+
557
+ def unregister_machine(
558
+ self, session_id: str, machine_id: Optional[str] = None
559
+ ) -> bool:
560
+ """Unregister a machine from a session."""
561
+ if not self.connection_manager.connected:
562
+ if not self.connect():
563
+ raise ConnectionError("Failed to connect to server")
564
+ return self.machine_manager.unregister_machine(session_id, machine_id)
565
+
566
+ def drain_machine(self, session_id: str, machine_id: Optional[str] = None) -> bool:
567
+ """Drain a machine from a session."""
568
+ if not self.connection_manager.connected:
569
+ if not self.connect():
570
+ raise ConnectionError("Failed to connect to server")
571
+ return self.machine_manager.drain_machine(session_id, machine_id)
572
+
573
+ def create_api_key(
574
+ self,
575
+ session_id: str,
576
+ name: str,
577
+ capabilities: List[str],
578
+ ) -> Dict[str, Any]:
579
+ """Create a new API key for a session.
580
+
581
+ capabilities is required and must be non-empty (least-privilege,
582
+ explicit key minting).
583
+ """
584
+ if not self.connection_manager.connected:
585
+ if not self.connect():
586
+ raise ConnectionError("Failed to connect to server")
587
+ return self.session_manager.create_api_key(session_id, name, capabilities)
588
+
589
+ def list_api_keys(self, session_id: str) -> List[Dict[str, Any]]:
590
+ """List active API keys for a session."""
591
+ if not self.connection_manager.connected:
592
+ if not self.connect():
593
+ raise ConnectionError("Failed to connect to server")
594
+ return self.session_manager.list_api_keys(session_id)
595
+
596
+ def revoke_api_key(self, session_id: str, key_id: str) -> bool:
597
+ """Revoke an API key for a session."""
598
+ if not self.connection_manager.connected:
599
+ if not self.connect():
600
+ raise ConnectionError("Failed to connect to server")
601
+ return self.session_manager.revoke_api_key(session_id, key_id)
602
+
603
+ def provider_runtime(self, session_ids: Optional[List[str]] = None):
604
+ """Return the explicit provider runtime for this client."""
605
+ from .provider_runtime import ProviderRuntime
606
+
607
+ if self._provider_runtime is None:
608
+ self._provider_runtime = ProviderRuntime(self, session_ids=session_ids)
609
+ elif session_ids:
610
+ self._provider_runtime.add_sessions(session_ids)
611
+ return self._provider_runtime
612
+
613
+ def start(self):
614
+ """Backward-compatible alias for the explicit provider runtime."""
615
+ self.provider_runtime().run_forever()
616
+
617
+ def stop(self):
618
+ """Stop the explicit provider runtime if it is active."""
619
+ self.running = False
620
+ if self._provider_runtime is not None:
621
+ self._provider_runtime.stop()
622
+
623
+ def _main_loop(self):
624
+ """Main polling loop for all sessions."""
625
+ while self.running:
626
+ try:
627
+ contexts = self.session_manager.list_session_contexts()
628
+
629
+ for context in contexts:
630
+ if context.machine_id:
631
+ context.poll_requests()
632
+
633
+ time.sleep(self.config.poll_interval)
634
+
635
+ except Exception as e:
636
+ logger.warning("Error in main loop: %s", e, exc_info=True)
637
+ time.sleep(1)
638
+
639
+ def __enter__(self):
640
+ """Context manager entry."""
641
+ self.connect()
642
+ return self
643
+
644
+ def __exit__(self, exc_type, exc_val, exc_tb):
645
+ """Context manager exit."""
646
+ self.disconnect()
647
+
648
+ # Legacy methods for backward compatibility
649
+ def health(self):
650
+ """Check server health."""
651
+ return self.connection_manager.health_check()
652
+
653
+ def _register_machine_internal(self):
654
+ """Legacy method for backward compatibility."""
655
+ return len(self.session_ids) > 0 and all(
656
+ self.get_session(session_id).machine_id is not None
657
+ for session_id in self.session_ids
658
+ )
659
+
660
+ def _register_tool(
661
+ self,
662
+ name: str,
663
+ func: Callable,
664
+ schema: Dict,
665
+ stream: bool = False,
666
+ tags: List[str] = None,
667
+ ):
668
+ """Legacy method for registering tools."""
669
+ # For backward compatibility, register on all sessions
670
+ for session_id in self.session_ids:
671
+ context = self.get_session(session_id)
672
+ if context:
673
+ context.register_tool(name, func, schema, stream=stream, tags=tags)
674
+
675
+ def _register_tools_with_server(self):
676
+ """Legacy method - tools are now registered immediately."""
677
+ pass
678
+
679
+ def _poll_pending_requests(self):
680
+ """Legacy method - polling is now handled by main loop."""
681
+ pass
@@ -0,0 +1,3 @@
1
+ from .schema import generate_schema_from_function, type_to_json_schema
2
+
3
+ __all__ = ["generate_schema_from_function", "type_to_json_schema"]