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.
- toolplane/__init__.py +106 -0
- toolplane/common/__init__.py +93 -0
- toolplane/common/base_config.py +129 -0
- toolplane/common/base_connection_manager.py +171 -0
- toolplane/common/base_session_manager.py +321 -0
- toolplane/common/base_tool_manager.py +347 -0
- toolplane/common/constants.py +47 -0
- toolplane/common/utils.py +310 -0
- toolplane/core/__init__.py +67 -0
- toolplane/core/config.py +107 -0
- toolplane/core/connection.py +285 -0
- toolplane/core/errors.py +298 -0
- toolplane/core/machine.py +480 -0
- toolplane/core/request.py +775 -0
- toolplane/core/session.py +332 -0
- toolplane/core/session_context.py +514 -0
- toolplane/core/task.py +130 -0
- toolplane/core/tool.py +329 -0
- toolplane/http_core/__init__.py +37 -0
- toolplane/http_core/http_config.py +97 -0
- toolplane/http_core/http_connection.py +409 -0
- toolplane/http_core/http_machine.py +298 -0
- toolplane/http_core/http_request.py +748 -0
- toolplane/http_core/http_session.py +348 -0
- toolplane/http_core/http_session_context.py +491 -0
- toolplane/http_core/http_task.py +101 -0
- toolplane/http_core/http_tool.py +400 -0
- toolplane/interfaces/__init__.py +27 -0
- toolplane/interfaces/client_interface.py +122 -0
- toolplane/interfaces/connection_interface.py +193 -0
- toolplane/interfaces/event_interface.py +290 -0
- toolplane/interfaces/request_interface.py +439 -0
- toolplane/interfaces/session_interface.py +288 -0
- toolplane/interfaces/tool_interface.py +441 -0
- toolplane/proto/__init__.py +0 -0
- toolplane/proto/service_pb2.py +315 -0
- toolplane/proto/service_pb2_grpc.py +2240 -0
- toolplane/provider_cli.py +268 -0
- toolplane/provider_registry.py +77 -0
- toolplane/provider_runtime.py +302 -0
- toolplane/toolkits/__init__.py +0 -0
- toolplane/toolkits/standalone_tools/__init__.py +0 -0
- toolplane/toolkits/standalone_tools/create_directory.py +94 -0
- toolplane/toolkits/standalone_tools/create_file.py +124 -0
- toolplane/toolkits/standalone_tools/file_search.py +229 -0
- toolplane/toolkits/standalone_tools/grep_search.py +372 -0
- toolplane/toolkits/standalone_tools/launcher.py +146 -0
- toolplane/toolkits/standalone_tools/list_dir.py +395 -0
- toolplane/toolkits/standalone_tools/read_file.py +346 -0
- toolplane/toolkits/standalone_tools/replace_string_in_file.py +407 -0
- toolplane/toolkits/standalone_tools/run_tests.py +66 -0
- toolplane/toolkits/standalone_tools/semantic_search.py +485 -0
- toolplane/toolkits/standalone_tools/standalone_toolkit.py +979 -0
- toolplane/toolkits/standalone_tools/test_failure_analysis.py +618 -0
- toolplane/toolkits/standalone_tools/test_standalone_toolkit.py +517 -0
- toolplane/toolkits/swe/__init__.py +35 -0
- toolplane/toolkits/swe/create_directory.py +15 -0
- toolplane/toolkits/swe/create_file.py +15 -0
- toolplane/toolkits/swe/descriptions.py +273 -0
- toolplane/toolkits/swe/execute_bash.py +93 -0
- toolplane/toolkits/swe/file_editor.py +775 -0
- toolplane/toolkits/swe/file_search.py +16 -0
- toolplane/toolkits/swe/finish.py +50 -0
- toolplane/toolkits/swe/grep_search.py +19 -0
- toolplane/toolkits/swe/list_dir.py +407 -0
- toolplane/toolkits/swe/read_file.py +18 -0
- toolplane/toolkits/swe/replace_string_in_file.py +17 -0
- toolplane/toolkits/swe/search.py +260 -0
- toolplane/toolkits/swe/semantic_search.py +20 -0
- toolplane/toolkits/swe/str_replace_editor.py +647 -0
- toolplane/toolkits/swe/submit.py +29 -0
- toolplane/toolkits/swe/swe_toolkit.py +1296 -0
- toolplane/toolplane_client.py +686 -0
- toolplane/toolplane_http_client.py +681 -0
- toolplane/utils/__init__.py +3 -0
- toolplane/utils/schema.py +146 -0
- toolplane_python_client-0.1.0.dist-info/METADATA +543 -0
- toolplane_python_client-0.1.0.dist-info/RECORD +81 -0
- toolplane_python_client-0.1.0.dist-info/WHEEL +5 -0
- toolplane_python_client-0.1.0.dist-info/entry_points.txt +2 -0
- 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
|