ioa-observe-sdk 1.0.29__py3-none-any.whl → 1.0.30__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 (6) hide show
  1. ioa_observe/sdk/instrumentations/slim.py +367 -844
  2. {ioa_observe_sdk-1.0.29.dist-info → ioa_observe_sdk-1.0.30.dist-info}/METADATA +1 -1
  3. {ioa_observe_sdk-1.0.29.dist-info → ioa_observe_sdk-1.0.30.dist-info}/RECORD +6 -6
  4. {ioa_observe_sdk-1.0.29.dist-info → ioa_observe_sdk-1.0.30.dist-info}/WHEEL +0 -0
  5. {ioa_observe_sdk-1.0.29.dist-info → ioa_observe_sdk-1.0.30.dist-info}/licenses/LICENSE.md +0 -0
  6. {ioa_observe_sdk-1.0.29.dist-info → ioa_observe_sdk-1.0.30.dist-info}/top_level.txt +0 -0
ioa_observe/sdk/instrumentations/slim.py CHANGED
@@ -1,6 +1,8 @@
1
1
  # Copyright AGNTCY Contributors (https://github.com/agntcy)
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
 
4
+ """SLIM v1.x Instrumentation for OpenTelemetry tracing."""
5
+
4
6
  from typing import Collection
5
7
  import functools
6
8
  import json
@@ -17,9 +19,121 @@ from ioa_observe.sdk import TracerWrapper
17
19
  from ioa_observe.sdk.client import kv_store
18
20
  from ioa_observe.sdk.tracing import set_session_id, get_current_traceparent
19
21
 
20
- _instruments = ("slim-bindings >= 0.4.0",)
22
+ _instruments = ("slim-bindings >= 1.0.0",)
21
23
  _global_tracer = None
22
- _kv_lock = threading.RLock() # Add thread-safety for kv_store operations
24
+ _kv_lock = threading.RLock()
25
+
26
+
27
+ def _get_session_id(session):
28
+ """Get session ID from session object (v1.x uses .session_id() method)."""
29
+ if hasattr(session, "session_id") and callable(session.session_id):
30
+ try:
31
+ return str(session.session_id())
32
+ except Exception:
33
+ pass
34
+ return str(session.id) if hasattr(session, "id") else None
35
+
36
+
37
+ def _process_received_message(raw_message):
38
+ """Process received message, extract tracing context, return cleaned payload."""
39
+ if raw_message is None:
40
+ return raw_message
41
+
42
+ try:
43
+ if isinstance(raw_message, bytes):
44
+ message_dict = json.loads(raw_message.decode())
45
+ elif isinstance(raw_message, str):
46
+ message_dict = json.loads(raw_message)
47
+ elif isinstance(raw_message, dict):
48
+ message_dict = raw_message
49
+ else:
50
+ return raw_message
51
+
52
+ headers = message_dict.get("headers", {})
53
+ traceparent = headers.get("traceparent")
54
+ session_id = headers.get("session_id")
55
+
56
+ # Restore trace context
57
+ if traceparent:
58
+ carrier = {
59
+ k.lower(): v
60
+ for k, v in headers.items()
61
+ if k.lower() in ["traceparent", "baggage"]
62
+ }
63
+ if carrier:
64
+ ctx = TraceContextTextMapPropagator().extract(carrier=carrier)
65
+ ctx = W3CBaggagePropagator().extract(carrier=carrier, context=ctx)
66
+ context.attach(ctx)
67
+
68
+ if session_id and session_id != "None":
69
+ set_session_id(session_id, traceparent=traceparent)
70
+ with _kv_lock:
71
+ kv_store.set(f"execution.{traceparent}", session_id)
72
+
73
+ # Clean headers
74
+ cleaned = message_dict.copy()
75
+ if "headers" in cleaned:
76
+ h = cleaned["headers"].copy()
77
+ for k in ["traceparent", "session_id", "slim_session_id"]:
78
+ h.pop(k, None)
79
+ if h:
80
+ cleaned["headers"] = h
81
+ else:
82
+ cleaned.pop("headers", None)
83
+
84
+ # Extract payload
85
+ if len(cleaned) == 1 and "payload" in cleaned:
86
+ payload = cleaned["payload"]
87
+ if isinstance(payload, str):
88
+ try:
89
+ return json.dumps(json.loads(payload)).encode("utf-8")
90
+ except json.JSONDecodeError:
91
+ return payload.encode("utf-8")
92
+ elif isinstance(payload, (dict, list)):
93
+ return json.dumps(payload).encode("utf-8")
94
+ return payload
95
+ return json.dumps(cleaned).encode("utf-8")
96
+
97
+ except Exception:
98
+ return raw_message
99
+
100
+
101
+ def _wrap_message_with_headers(message, headers):
102
+ """Wrap message with tracing headers."""
103
+ if isinstance(message, bytes):
104
+ try:
105
+ decoded = message.decode("utf-8")
106
+ try:
107
+ original = json.loads(decoded)
108
+ if isinstance(original, dict):
109
+ wrapped = original.copy()
110
+ wrapped["headers"] = {**wrapped.get("headers", {}), **headers}
111
+ else:
112
+ wrapped = {"headers": headers, "payload": original}
113
+ except json.JSONDecodeError:
114
+ wrapped = {"headers": headers, "payload": decoded}
115
+ except UnicodeDecodeError:
116
+ wrapped = {
117
+ "headers": headers,
118
+ "payload": base64.b64encode(message).decode("utf-8"),
119
+ }
120
+ elif isinstance(message, str):
121
+ try:
122
+ original = json.loads(message)
123
+ if isinstance(original, dict):
124
+ wrapped = original.copy()
125
+ wrapped["headers"] = {**wrapped.get("headers", {}), **headers}
126
+ else:
127
+ wrapped = {"headers": headers, "payload": original}
128
+ except json.JSONDecodeError:
129
+ wrapped = {"headers": headers, "payload": message}
130
+ elif isinstance(message, dict):
131
+ wrapped = message.copy()
132
+ wrapped["headers"] = {**wrapped.get("headers", {}), **headers}
133
+ else:
134
+ wrapped = {"headers": headers, "payload": json.dumps(message)}
135
+
136
+ return wrapped
23
137
 
24
138
 
25
139
  class SLIMInstrumentor(BaseInstrumentor):
@@ -39,935 +153,344 @@ class SLIMInstrumentor(BaseInstrumentor):
39
153
  "No module named 'slim_bindings'. Please install it first."
40
154
  )
41
155
 
42
- # Instrument `publish` method - handles multiple signatures
43
- # In v0.6.0+, publish moved from Slim class to Session objects
44
- if hasattr(slim_bindings.Slim, "publish"):
45
- # Legacy v0.5.x app-level publish method
46
- original_publish = slim_bindings.Slim.publish
47
-
48
- @functools.wraps(original_publish)
49
- async def instrumented_publish(self, *args, **kwargs):
50
- if _global_tracer:
51
- with _global_tracer.start_as_current_span("slim.publish") as span:
52
- traceparent = get_current_traceparent()
53
-
54
- # Handle different publish signatures
55
- # Definition 1: publish(session, message, topic_name) - v0.4.0+ group chat
56
- # Definition 2: publish(session, message, organization, namespace, topic) - legacy
57
- if len(args) >= 3:
58
- session_arg = args[0] if args else None
59
- if hasattr(session_arg, "id"):
60
- span.set_attribute(
61
- "slim.session.id", str(session_arg.id)
62
- )
63
-
64
- # Check if third argument is PyName (new API) or string (legacy API)
65
- if len(args) >= 3 and hasattr(args[2], "organization"):
66
- # New API: args[2] is PyName
67
- topic_name = args[2]
68
- span.set_attribute(
69
- "slim.topic.organization", topic_name.organization
70
- )
71
- span.set_attribute(
72
- "slim.topic.namespace", topic_name.namespace
73
- )
74
- span.set_attribute("slim.topic.app", topic_name.app)
75
- else:
76
- traceparent = get_current_traceparent()
77
-
78
- # Thread-safe access to kv_store
79
- session_id = None
80
- if traceparent:
81
- with _kv_lock:
82
- session_id = kv_store.get(f"execution.{traceparent}")
83
-
84
- headers = {
85
- "session_id": session_id if session_id else None,
86
- "traceparent": traceparent,
87
- }
88
-
89
- # Set baggage context
90
- if traceparent and session_id:
91
- baggage.set_baggage(f"execution.{traceparent}", session_id)
92
-
93
- # Wrap message with headers - handle different message positions
94
- message_arg_index = 1 # message will typically be the second argument
95
- if len(args) > message_arg_index:
96
- original_args = list(args)
97
- message = original_args[message_arg_index]
98
- wrapped_message = SLIMInstrumentor._wrap_message_with_headers(
99
- self, message, headers
100
- )
101
-
102
- # Convert wrapped message back to bytes if needed
103
- if isinstance(wrapped_message, dict):
104
- message_to_send = json.dumps(wrapped_message).encode("utf-8")
105
- else:
106
- message_to_send = wrapped_message
156
+ self._instrument_service(slim_bindings)
157
+ self._instrument_app(slim_bindings)
158
+ self._instrument_sessions(slim_bindings)
107
159
 
108
- original_args[message_arg_index] = message_to_send
109
- args = tuple(original_args)
160
+ def _instrument_service(self, slim_bindings):
161
+ """Instrument SLIM v1.x Service class."""
162
+ if not hasattr(slim_bindings, "Service"):
163
+ return
110
164
 
111
- return await original_publish(self, *args, **kwargs)
165
+ Service = slim_bindings.Service
112
166
 
113
- slim_bindings.Slim.publish = instrumented_publish
167
+ # connect_async
168
+ if hasattr(Service, "connect_async"):
169
+ original_connect_async = Service.connect_async
114
170
 
115
- # Instrument `publish_to` (new v0.4.0+ method)
116
- if hasattr(slim_bindings.Slim, "publish_to"):
117
- original_publish_to = slim_bindings.Slim.publish_to
118
-
119
- @functools.wraps(original_publish_to)
120
- async def instrumented_publish_to(
121
- self, session_info, message, *args, **kwargs
122
- ):
171
+ @functools.wraps(original_connect_async)
172
+ async def wrapped_connect_async(self, config, *args, **kwargs):
123
173
  if _global_tracer:
124
174
  with _global_tracer.start_as_current_span(
125
- "slim.publish_to"
175
+ "slim.service.connect"
126
176
  ) as span:
127
- traceparent = get_current_traceparent()
177
+ result = await original_connect_async(
178
+ self, config, *args, **kwargs
179
+ )
180
+ span.set_attribute(
181
+ "slim.connection.id", str(result) if result else "unknown"
182
+ )
183
+ return result
184
+ return await original_connect_async(self, config, *args, **kwargs)
128
185
 
129
- # Add session context to span
130
- if hasattr(session_info, "id"):
131
- span.set_attribute("slim.session.id", str(session_info.id))
132
- else:
133
- traceparent = get_current_traceparent()
186
+ Service.connect_async = wrapped_connect_async
134
187
 
135
- # Thread-safe access to kv_store
136
- session_id = None
137
- if traceparent:
138
- with _kv_lock:
139
- session_id = kv_store.get(f"execution.{traceparent}")
140
- if session_id:
141
- kv_store.set(f"execution.{traceparent}", session_id)
188
+ # run_server_async
189
+ if hasattr(Service, "run_server_async"):
190
+ original_run_server_async = Service.run_server_async
142
191
 
143
- headers = {
144
- "session_id": session_id if session_id else None,
145
- "traceparent": traceparent,
146
- "slim_session_id": str(session_info.id)
147
- if hasattr(session_info, "id")
148
- else None,
149
- }
150
-
151
- # Set baggage context
152
- if traceparent and session_id:
153
- baggage.set_baggage(f"execution.{traceparent}", session_id)
192
+ @functools.wraps(original_run_server_async)
193
+ async def wrapped_run_server_async(self, config, *args, **kwargs):
194
+ if _global_tracer:
195
+ with _global_tracer.start_as_current_span(
196
+ "slim.service.run_server"
197
+ ) as span:
198
+ if hasattr(config, "endpoint"):
199
+ span.set_attribute("slim.server.endpoint", config.endpoint)
200
+ result = await original_run_server_async(
201
+ self, config, *args, **kwargs
202
+ )
203
+ return result
204
+ return await original_run_server_async(self, config, *args, **kwargs)
154
205
 
155
- wrapped_message = SLIMInstrumentor._wrap_message_with_headers(
156
- self, message, headers
157
- )
158
- message_to_send = (
159
- json.dumps(wrapped_message).encode("utf-8")
160
- if isinstance(wrapped_message, dict)
161
- else wrapped_message
162
- )
206
+ Service.run_server_async = wrapped_run_server_async
163
207
 
164
- return await original_publish_to(
165
- self, session_info, message_to_send, *args, **kwargs
166
- )
208
+ def _instrument_app(self, slim_bindings):
209
+ """Instrument SLIM v1.x App class."""
210
+ if not hasattr(slim_bindings, "App"):
211
+ return
167
212
 
168
- slim_bindings.Slim.publish_to = instrumented_publish_to
213
+ App = slim_bindings.App
169
214
 
170
- # Instrument `request_reply` (v0.4.0+ to v0.5.x method, removed in v0.6.0)
171
- if hasattr(slim_bindings.Slim, "request_reply"):
172
- original_request_reply = slim_bindings.Slim.request_reply
215
+ # create_session_async
216
+ if hasattr(App, "create_session_async"):
217
+ original_create_session_async = App.create_session_async
173
218
 
174
- @functools.wraps(original_request_reply)
175
- async def instrumented_request_reply(
176
- self, session_info, message, remote_name, timeout=None, *args, **kwargs
219
+ @functools.wraps(original_create_session_async)
220
+ async def wrapped_create_session_async(
221
+ self, config, dest=None, *args, **kwargs
177
222
  ):
178
223
  if _global_tracer:
179
224
  with _global_tracer.start_as_current_span(
180
- "slim.request_reply"
225
+ "slim.app.create_session"
181
226
  ) as span:
182
- traceparent = get_current_traceparent()
183
-
184
- # Add context to span
185
- if hasattr(session_info, "id"):
186
- span.set_attribute("slim.session.id", str(session_info.id))
187
- if hasattr(remote_name, "organization"):
188
- span.set_attribute(
189
- "slim.remote.organization", remote_name.organization
190
- )
191
- span.set_attribute(
192
- "slim.remote.namespace", remote_name.namespace
193
- )
194
- span.set_attribute("slim.remote.app", remote_name.app)
195
- else:
196
- traceparent = get_current_traceparent()
197
-
198
- # Thread-safe access to kv_store
199
- session_id = None
200
- if traceparent:
201
- with _kv_lock:
202
- session_id = kv_store.get(f"execution.{traceparent}")
203
-
204
- headers = {
205
- "session_id": session_id if session_id else None,
206
- "traceparent": traceparent,
207
- "slim_session_id": str(session_info.id)
208
- if hasattr(session_info, "id")
209
- else None,
210
- }
211
-
212
- # Set baggage context
213
- if traceparent and session_id:
214
- baggage.set_baggage(f"execution.{traceparent}", session_id)
215
-
216
- wrapped_message = SLIMInstrumentor._wrap_message_with_headers(
217
- self, message, headers
218
- )
219
- message_to_send = (
220
- json.dumps(wrapped_message).encode("utf-8")
221
- if isinstance(wrapped_message, dict)
222
- else wrapped_message
223
- )
224
-
225
- kwargs_with_timeout = kwargs.copy()
226
- if timeout is not None:
227
- kwargs_with_timeout["timeout"] = timeout
228
-
229
- return await original_request_reply(
230
- self,
231
- session_info,
232
- message_to_send,
233
- remote_name,
234
- **kwargs_with_timeout,
227
+ ctx = await original_create_session_async(
228
+ self, config, dest, *args, **kwargs
229
+ )
230
+ if hasattr(ctx, "session"):
231
+ sid = _get_session_id(ctx.session)
232
+ if sid:
233
+ span.set_attribute("slim.session.id", sid)
234
+ return ctx
235
+ return await original_create_session_async(
236
+ self, config, dest, *args, **kwargs
235
237
  )
236
238
 
237
- slim_bindings.Slim.request_reply = instrumented_request_reply
239
+ App.create_session_async = wrapped_create_session_async
238
240
 
239
- # Instrument `invite` (new v0.4.0+ method for group chat)
240
- if hasattr(slim_bindings.Slim, "invite"):
241
- original_invite = slim_bindings.Slim.invite
241
+ # create_session_and_wait_async
242
+ if hasattr(App, "create_session_and_wait_async"):
243
+ original_create_session_and_wait_async = App.create_session_and_wait_async
242
244
 
243
- @functools.wraps(original_invite)
244
- async def instrumented_invite(
245
- self, session_info, participant_name, *args, **kwargs
245
+ @functools.wraps(original_create_session_and_wait_async)
246
+ async def wrapped_create_session_and_wait_async(
247
+ self, config, dest=None, *args, **kwargs
246
248
  ):
247
249
  if _global_tracer:
248
- with _global_tracer.start_as_current_span("slim.invite") as span:
249
- # Add context to span
250
- if hasattr(session_info, "id"):
251
- span.set_attribute("slim.session.id", str(session_info.id))
252
- if hasattr(participant_name, "organization"):
253
- span.set_attribute(
254
- "slim.participant.organization",
255
- participant_name.organization,
256
- )
257
- span.set_attribute(
258
- "slim.participant.namespace", participant_name.namespace
259
- )
260
- span.set_attribute(
261
- "slim.participant.app", participant_name.app
262
- )
263
-
264
- return await original_invite(
265
- self, session_info, participant_name, *args, **kwargs
250
+ with _global_tracer.start_as_current_span(
251
+ "slim.app.create_session"
252
+ ) as span:
253
+ session = await original_create_session_and_wait_async(
254
+ self, config, dest, *args, **kwargs
255
+ )
256
+ sid = _get_session_id(session)
257
+ if sid:
258
+ span.set_attribute("slim.session.id", sid)
259
+ return session
260
+ return await original_create_session_and_wait_async(
261
+ self, config, dest, *args, **kwargs
266
262
  )
267
263
 
268
- slim_bindings.Slim.invite = instrumented_invite
264
+ App.create_session_and_wait_async = wrapped_create_session_and_wait_async
269
265
 
270
- # Instrument `set_route` (new v0.4.0+ method)
271
- if hasattr(slim_bindings.Slim, "set_route"):
272
- original_set_route = slim_bindings.Slim.set_route
266
+ # subscribe_async
267
+ if hasattr(App, "subscribe_async"):
268
+ original_subscribe_async = App.subscribe_async
273
269
 
274
- @functools.wraps(original_set_route)
275
- async def instrumented_set_route(self, remote_name, *args, **kwargs):
270
+ @functools.wraps(original_subscribe_async)
271
+ async def wrapped_subscribe_async(self, name, conn_id, *args, **kwargs):
276
272
  if _global_tracer:
277
- with _global_tracer.start_as_current_span("slim.set_route") as span:
278
- # Add context to span
279
- if hasattr(remote_name, "organization"):
280
- span.set_attribute(
281
- "slim.route.organization", remote_name.organization
282
- )
273
+ with _global_tracer.start_as_current_span(
274
+ "slim.app.subscribe"
275
+ ) as span:
276
+ if hasattr(name, "organization"):
283
277
  span.set_attribute(
284
- "slim.route.namespace", remote_name.namespace
278
+ "slim.name",
279
+ f"{name.organization}/{name.namespace}/{name.app}",
285
280
  )
286
- span.set_attribute("slim.route.app", remote_name.app)
287
-
288
- return await original_set_route(self, remote_name, *args, **kwargs)
289
-
290
- slim_bindings.Slim.set_route = instrumented_set_route
291
-
292
- # Instrument `receive` - only if it exists (removed in v0.6.0)
293
- if hasattr(slim_bindings.Slim, "receive"):
294
- original_receive = slim_bindings.Slim.receive
295
-
296
- @functools.wraps(original_receive)
297
- async def instrumented_receive(
298
- self, session=None, timeout=None, *args, **kwargs
299
- ):
300
- # Handle both old and new API patterns
301
- if session is not None or timeout is not None:
302
- # New API pattern with session parameter
303
- kwargs_with_params = kwargs.copy()
304
- if session is not None:
305
- kwargs_with_params["session"] = session
306
- if timeout is not None:
307
- kwargs_with_params["timeout"] = timeout
308
- recv_session, raw_message = await original_receive(
309
- self, **kwargs_with_params
310
- )
311
- else:
312
- # Legacy API pattern
313
- recv_session, raw_message = await original_receive(
314
- self, *args, **kwargs
315
- )
316
-
317
- if raw_message is None:
318
- return recv_session, raw_message
319
-
320
- try:
321
- message_dict = json.loads(raw_message.decode())
322
- headers = message_dict.get("headers", {})
323
-
324
- # Extract traceparent and session info from headers
325
- traceparent = headers.get("traceparent")
326
- session_id = headers.get("session_id")
327
-
328
- # Create carrier for context propagation
329
- carrier = {}
330
- for key in ["traceparent", "Traceparent", "baggage", "Baggage"]:
331
- if key.lower() in [k.lower() for k in headers.keys()]:
332
- for k in headers.keys():
333
- if k.lower() == key.lower():
334
- carrier[key.lower()] = headers[k]
335
-
336
- # Restore trace context
337
- if carrier and traceparent:
338
- ctx = TraceContextTextMapPropagator().extract(carrier=carrier)
339
- ctx = W3CBaggagePropagator().extract(
340
- carrier=carrier, context=ctx
341
- )
342
-
343
- # Activate the restored context
344
- token = context.attach(ctx)
345
-
346
- try:
347
- # Set execution ID with the restored context
348
- if session_id and session_id != "None":
349
- set_session_id(session_id, traceparent=traceparent)
350
-
351
- # Store in kv_store with thread safety
352
- with _kv_lock:
353
- kv_store.set(f"execution.{traceparent}", session_id)
354
-
355
- # DON'T detach the context yet - we need it to persist for the callback
356
- # The context will be cleaned up later or by the garbage collector
357
-
358
- except Exception as e:
359
- # Only detach on error
360
- context.detach(token)
361
- raise e
362
- elif traceparent and session_id and session_id != "None":
363
- # Even without carrier context, set session ID if we have the data
364
- set_session_id(session_id, traceparent=traceparent)
365
-
366
- # Fallback: check stored execution ID if not found in headers
367
- if traceparent and (not session_id or session_id == "None"):
368
- with _kv_lock:
369
- stored_session_id = kv_store.get(f"execution.{traceparent}")
370
- if stored_session_id:
371
- session_id = stored_session_id
372
- set_session_id(session_id, traceparent=traceparent)
373
-
374
- # Process and clean the message
375
- message_to_return = message_dict.copy()
376
- if "headers" in message_to_return:
377
- headers_copy = message_to_return["headers"].copy()
378
- # Remove tracing-specific headers but keep other headers
379
- headers_copy.pop("traceparent", None)
380
- headers_copy.pop("session_id", None)
381
- headers_copy.pop("slim_session_id", None)
382
- if headers_copy:
383
- message_to_return["headers"] = headers_copy
384
- else:
385
- message_to_return.pop("headers", None)
386
-
387
- # Return processed message
388
- if len(message_to_return) == 1 and "payload" in message_to_return:
389
- payload = message_to_return["payload"]
390
- if isinstance(payload, str):
391
- try:
392
- payload_dict = json.loads(payload)
393
- return recv_session, json.dumps(payload_dict).encode(
394
- "utf-8"
395
- )
396
- except json.JSONDecodeError:
397
- return recv_session, payload.encode(
398
- "utf-8"
399
- ) if isinstance(payload, str) else payload
400
- return recv_session, json.dumps(payload).encode(
401
- "utf-8"
402
- ) if isinstance(payload, (dict, list)) else payload
403
- else:
404
- return recv_session, json.dumps(message_to_return).encode(
405
- "utf-8"
281
+ return await original_subscribe_async(
282
+ self, name, conn_id, *args, **kwargs
406
283
  )
284
+ return await original_subscribe_async(
285
+ self, name, conn_id, *args, **kwargs
286
+ )
407
287
 
408
- except Exception as e:
409
- print(f"Error processing message: {e}")
410
- return recv_session, raw_message
411
-
412
- slim_bindings.Slim.receive = instrumented_receive
413
-
414
- # Instrument `connect` - only if it exists
415
- if hasattr(slim_bindings.Slim, "connect"):
416
- original_connect = slim_bindings.Slim.connect
417
-
418
- @functools.wraps(original_connect)
419
- async def instrumented_connect(self, *args, **kwargs):
420
- if _global_tracer:
421
- with _global_tracer.start_as_current_span("slim.connect"):
422
- return await original_connect(self, *args, **kwargs)
423
- else:
424
- return await original_connect(self, *args, **kwargs)
425
-
426
- slim_bindings.Slim.connect = instrumented_connect
288
+ App.subscribe_async = wrapped_subscribe_async
427
289
 
428
- # Instrument `create_session` (new v0.4.0+ method)
429
- if hasattr(slim_bindings.Slim, "create_session"):
430
- original_create_session = slim_bindings.Slim.create_session
290
+ # set_route_async
291
+ if hasattr(App, "set_route_async"):
292
+ original_set_route_async = App.set_route_async
431
293
 
432
- @functools.wraps(original_create_session)
433
- async def instrumented_create_session(self, config, *args, **kwargs):
294
+ @functools.wraps(original_set_route_async)
295
+ async def wrapped_set_route_async(self, name, conn_id, *args, **kwargs):
434
296
  if _global_tracer:
435
297
  with _global_tracer.start_as_current_span(
436
- "slim.create_session"
298
+ "slim.app.set_route"
437
299
  ) as span:
438
- session_info = await original_create_session(
439
- self, config, *args, **kwargs
300
+ if hasattr(name, "organization"):
301
+ span.set_attribute(
302
+ "slim.route",
303
+ f"{name.organization}/{name.namespace}/{name.app}",
304
+ )
305
+ return await original_set_route_async(
306
+ self, name, conn_id, *args, **kwargs
440
307
  )
308
+ return await original_set_route_async(
309
+ self, name, conn_id, *args, **kwargs
310
+ )
441
311
 
442
- # Add session attributes to span
443
- if hasattr(session_info, "id"):
444
- span.set_attribute("slim.session.id", str(session_info.id))
445
-
446
- return session_info
447
- else:
448
- return await original_create_session(self, config, *args, **kwargs)
449
-
450
- slim_bindings.Slim.create_session = instrumented_create_session
451
-
452
- # Instrument new v0.6.0+ session-level methods
453
- # These methods are available on Session objects, not the Slim app
454
- self._instrument_session_methods(slim_bindings)
312
+ App.set_route_async = wrapped_set_route_async
455
313
 
456
- # Instrument new v0.6.0+ app-level methods
457
- # listen_for_session replaces app.receive() for new sessions in v0.6.0+
458
- if hasattr(slim_bindings.Slim, "listen_for_session"):
459
- original_listen_for_session = slim_bindings.Slim.listen_for_session
314
+ # listen_for_session - this is a blocking call, not async
315
+ if hasattr(App, "listen_for_session"):
316
+ original_listen_for_session = App.listen_for_session
460
317
 
461
318
  @functools.wraps(original_listen_for_session)
462
- async def instrumented_listen_for_session(self, *args, **kwargs):
319
+ def wrapped_listen_for_session(self, *args, **kwargs):
463
320
  if _global_tracer:
464
321
  with _global_tracer.start_as_current_span(
465
- "slim.listen_for_session"
322
+ "slim.app.listen_for_session"
466
323
  ):
467
- session = await original_listen_for_session(
468
- self, *args, **kwargs
469
- )
324
+ return original_listen_for_session(self, *args, **kwargs)
325
+ return original_listen_for_session(self, *args, **kwargs)
470
326
 
471
- return session
472
- else:
473
- return await original_listen_for_session(self, *args, **kwargs)
474
-
475
- slim_bindings.Slim.listen_for_session = instrumented_listen_for_session
476
-
477
- def _instrument_session_methods(self, slim_bindings):
478
- # Try to find session-related classes in the slim_bindings module
479
- session_classes = []
480
-
481
- # Check for v0.6.0+ Session classes
482
- if hasattr(slim_bindings, "Session"):
483
- for attr_name in ["Session", "P2PSession", "GroupSession"]:
484
- if hasattr(slim_bindings, attr_name):
485
- session_class = getattr(slim_bindings, attr_name)
486
- session_classes.append((attr_name, session_class))
487
-
488
- # Check for older PySession class (pre-v0.6.0)
489
- if hasattr(slim_bindings, "PySession"):
490
- session_classes.append(("PySession", slim_bindings.PySession))
491
-
492
- # Also look for any class that has session-like methods
493
- for attr_name in dir(slim_bindings):
494
- attr = getattr(slim_bindings, attr_name)
495
- if isinstance(attr, type) and (
496
- hasattr(attr, "get_message") or hasattr(attr, "publish")
497
- ):
498
- session_classes.append((attr_name, attr))
327
+ App.listen_for_session = wrapped_listen_for_session
328
+
329
+ def _instrument_sessions(self, slim_bindings):
330
+ """Instrument session classes for v1.x."""
331
+ session_classes = set()
499
332
 
500
- # Instrument session methods for found classes
501
- for class_name, session_class in session_classes:
502
- # Instrument get_message (v0.6.0+ replacement for receive)
503
- if hasattr(session_class, "get_message"):
504
- self._instrument_session_get_message(session_class)
333
+ # Find session classes
334
+ for name in ["Session", "P2PSession", "GroupSession"]:
335
+ if hasattr(slim_bindings, name):
336
+ session_classes.add(getattr(slim_bindings, name))
505
337
 
506
- # Instrument session publish methods
507
- if hasattr(session_class, "publish"):
508
- self._instrument_session_publish(session_class, "publish")
338
+ # Find any class with session-like methods
339
+ for name in dir(slim_bindings):
340
+ cls = getattr(slim_bindings, name)
341
+ if isinstance(cls, type) and any(
342
+ hasattr(cls, m) for m in ["get_message_async", "publish_async"]
343
+ ):
344
+ session_classes.add(cls)
509
345
 
510
- if hasattr(session_class, "publish_to"):
511
- self._instrument_session_publish(session_class, "publish_to")
346
+ for session_class in session_classes:
347
+ if hasattr(session_class, "get_message_async"):
348
+ self._wrap_get_message(session_class)
512
349
 
513
- def _instrument_session_get_message(self, session_class):
514
- """Instrument session.get_message method (v0.6.0+)"""
515
- original_get_message = session_class.get_message
350
+ if hasattr(session_class, "publish_async"):
351
+ self._wrap_publish(session_class, "publish_async")
516
352
 
517
- @functools.wraps(original_get_message)
518
- async def instrumented_get_message(self, timeout=None, *args, **kwargs):
519
- # Handle the message reception similar to the old receive method
520
- if timeout is not None:
521
- kwargs["timeout"] = timeout
353
+ if hasattr(session_class, "publish_to_async"):
354
+ self._wrap_publish(session_class, "publish_to_async", msg_idx=1)
522
355
 
523
- result = await original_get_message(self, **kwargs)
356
+ def _wrap_get_message(self, session_class):
357
+ """Wrap get_message_async to extract tracing context."""
358
+ orig = session_class.get_message_async
524
359
 
525
- # Handle different return types from get_message
360
+ @functools.wraps(orig)
361
+ async def wrapped(self, *args, **kwargs):
362
+ result = await orig(self, *args, **kwargs)
526
363
  if result is None:
527
364
  return result
528
365
 
529
- # Check if get_message returns a tuple (context, message) or just message
530
- if isinstance(result, tuple) and len(result) == 2:
531
- message_context, raw_message = result
366
+ # v1.x returns ReceivedMessage with .context and .payload
367
+ if hasattr(result, "payload"):
368
+ raw = result.payload
369
+ elif isinstance(result, tuple) and len(result) == 2:
370
+ raw = result[1]
532
371
  else:
533
- raw_message = result
534
- message_context = None
372
+ raw = result
535
373
 
536
- if raw_message is None:
537
- return result
374
+ processed = _process_received_message(raw)
538
375
 
539
- try:
540
- # Handle different message types
541
- if isinstance(raw_message, bytes):
542
- message_dict = json.loads(raw_message.decode())
543
- elif isinstance(raw_message, str):
544
- message_dict = json.loads(raw_message)
545
- elif isinstance(raw_message, dict):
546
- message_dict = raw_message
547
- else:
548
- # Unknown type, return as-is
549
- return result
550
-
551
- headers = message_dict.get("headers", {})
552
-
553
- # Extract traceparent and session info from headers
554
- traceparent = headers.get("traceparent")
555
- session_id = headers.get("session_id")
556
-
557
- # Create carrier for context propagation
558
- carrier = {}
559
- for key in ["traceparent", "Traceparent", "baggage", "Baggage"]:
560
- if key.lower() in [k.lower() for k in headers.keys()]:
561
- for k in headers.keys():
562
- if k.lower() == key.lower():
563
- carrier[key.lower()] = headers[k]
564
-
565
- # Restore trace context
566
- if carrier and traceparent:
567
- ctx = TraceContextTextMapPropagator().extract(carrier=carrier)
568
- ctx = W3CBaggagePropagator().extract(carrier=carrier, context=ctx)
569
-
570
- # Activate the restored context
571
- token = context.attach(ctx)
572
-
573
- try:
574
- # Set execution ID with the restored context
575
- if session_id and session_id != "None":
576
- set_session_id(session_id, traceparent=traceparent)
577
-
578
- # Store in kv_store with thread safety
579
- with _kv_lock:
580
- kv_store.set(f"execution.{traceparent}", session_id)
581
-
582
- # DON'T detach the context yet - we need it to persist for the callback
583
-
584
- except Exception as e:
585
- # Only detach on error
586
- context.detach(token)
587
- raise e
588
- elif traceparent and session_id and session_id != "None":
589
- # Even without carrier context, set session ID if we have the data
590
- set_session_id(session_id, traceparent=traceparent)
591
-
592
- # Fallback: check stored execution ID if not found in headers
593
- if traceparent and (not session_id or session_id == "None"):
594
- with _kv_lock:
595
- stored_session_id = kv_store.get(f"execution.{traceparent}")
596
- if stored_session_id:
597
- session_id = stored_session_id
598
- set_session_id(session_id, traceparent=traceparent)
599
-
600
- # Process and clean the message
601
- message_to_return = message_dict.copy()
602
- if "headers" in message_to_return:
603
- headers_copy = message_to_return["headers"].copy()
604
- # Remove tracing-specific headers but keep other headers
605
- headers_copy.pop("traceparent", None)
606
- headers_copy.pop("session_id", None)
607
- headers_copy.pop("slim_session_id", None)
608
- if headers_copy:
609
- message_to_return["headers"] = headers_copy
610
- else:
611
- message_to_return.pop("headers", None)
612
-
613
- # Return processed message, maintaining original return format
614
- if len(message_to_return) == 1 and "payload" in message_to_return:
615
- payload = message_to_return["payload"]
616
- if isinstance(payload, str):
617
- try:
618
- payload_dict = json.loads(payload)
619
- processed_message = json.dumps(payload_dict).encode("utf-8")
620
- except json.JSONDecodeError:
621
- processed_message = (
622
- payload.encode("utf-8")
623
- if isinstance(payload, str)
624
- else payload
625
- )
626
- else:
627
- processed_message = (
628
- json.dumps(payload).encode("utf-8")
629
- if isinstance(payload, (dict, list))
630
- else payload
631
- )
632
- else:
633
- processed_message = json.dumps(message_to_return).encode("utf-8")
376
+ if hasattr(result, "payload"):
377
+ try:
378
+ result.payload = processed
379
+ except AttributeError:
634
380
 
635
- # Return in the same format as received
636
- if isinstance(result, tuple) and len(result) == 2:
637
- return (message_context, processed_message)
638
- else:
639
- return processed_message
381
+ class Processed:
382
+ def __init__(self, ctx, payload):
383
+ self.context = ctx
384
+ self.payload = payload
640
385
 
641
- except Exception as e:
642
- print(f"Error processing message: {e}")
386
+ return Processed(result.context, processed)
643
387
  return result
388
+ elif isinstance(result, tuple) and len(result) == 2:
389
+ return (result[0], processed)
390
+ return processed
391
+
392
+ session_class.get_message_async = wrapped
393
+
394
+ def _wrap_publish(self, session_class, method_name, msg_idx=0):
395
+ """Wrap publish methods to inject tracing headers."""
396
+ orig = getattr(session_class, method_name)
397
+
398
+ @functools.wraps(orig)
399
+ async def wrapped(self, *args, **kwargs):
400
+ traceparent = get_current_traceparent()
401
+ session_id = None
402
+
403
+ if traceparent:
404
+ with _kv_lock:
405
+ session_id = kv_store.get(f"execution.{traceparent}")
406
+ if not session_id:
407
+ session_id = get_value("session.id")
408
+ if session_id:
409
+ kv_store.set(f"execution.{traceparent}", session_id)
644
410
 
645
- session_class.get_message = instrumented_get_message
646
-
647
- def _instrument_session_publish(self, session_class, method_name):
648
- """Instrument session publish methods"""
649
- original_method = getattr(session_class, method_name)
411
+ slim_session_id = _get_session_id(self)
650
412
 
651
- @functools.wraps(original_method)
652
- async def instrumented_session_publish(self, *args, **kwargs):
653
413
  if _global_tracer:
654
414
  with _global_tracer.start_as_current_span(
655
415
  f"session.{method_name}"
656
416
  ) as span:
657
- traceparent = get_current_traceparent()
658
-
659
- # Add session context to span
660
- if hasattr(self, "id"):
661
- span.set_attribute("slim.session.id", str(self.id))
662
-
663
- # Handle message wrapping
664
- if args:
665
- # Thread-safe access to kv_store
666
- session_id = None
667
- if traceparent:
668
- with _kv_lock:
669
- session_id = kv_store.get(f"execution.{traceparent}")
670
- if not session_id:
671
- session_id = get_value("session.id")
672
- if session_id:
673
- kv_store.set(
674
- f"execution.{traceparent}", session_id
675
- )
417
+ if slim_session_id:
418
+ span.set_attribute("slim.session.id", slim_session_id)
676
419
 
420
+ if args and len(args) > msg_idx and (traceparent or session_id):
677
421
  headers = {
678
- "session_id": session_id if session_id else None,
422
+ "session_id": session_id,
679
423
  "traceparent": traceparent,
680
- "slim_session_id": str(self.id)
681
- if hasattr(self, "id")
682
- else None,
424
+ "slim_session_id": slim_session_id,
683
425
  }
684
-
685
- # Set baggage context
686
426
  if traceparent and session_id:
687
427
  baggage.set_baggage(f"execution.{traceparent}", session_id)
688
428
 
689
- # Wrap the message (first argument for publish, second for publish_to)
690
- # If the session_class is SessionContext, the message is always in second position
691
- message_idx = (
692
- 1
693
- if method_name == "publish_to"
694
- or session_class.__name__ == "SessionContext"
695
- else 0
429
+ args_list = list(args)
430
+ wrapped_msg = _wrap_message_with_headers(
431
+ args_list[msg_idx], headers
696
432
  )
697
- if len(args) > message_idx:
698
- args_list = list(args)
699
- message = args_list[message_idx]
700
- wrapped_message = (
701
- SLIMInstrumentor._wrap_message_with_headers(
702
- None, message, headers
703
- )
704
- )
705
-
706
- # Convert wrapped message back to bytes if needed
707
- if isinstance(wrapped_message, dict):
708
- message_to_send = json.dumps(wrapped_message).encode(
709
- "utf-8"
710
- )
711
- else:
712
- message_to_send = wrapped_message
713
-
714
- args_list[message_idx] = message_to_send
715
- args = tuple(args_list)
433
+ args_list[msg_idx] = (
434
+ json.dumps(wrapped_msg).encode("utf-8")
435
+ if isinstance(wrapped_msg, dict)
436
+ else wrapped_msg
437
+ )
438
+ args = tuple(args_list)
716
439
 
717
- return await original_method(self, *args, **kwargs)
440
+ return await orig(self, *args, **kwargs)
718
441
  else:
719
- # Handle message wrapping even without tracing
720
- if args:
721
- traceparent = get_current_traceparent()
722
- session_id = None
723
- if traceparent:
724
- with _kv_lock:
725
- session_id = kv_store.get(f"execution.{traceparent}")
726
- if not session_id:
727
- session_id = get_value("session.id")
728
- if session_id:
729
- kv_store.set(f"execution.{traceparent}", session_id)
730
-
731
- if traceparent or session_id:
732
- headers = {
733
- "session_id": session_id if session_id else None,
734
- "traceparent": traceparent,
735
- "slim_session_id": str(self.id)
736
- if hasattr(self, "id")
737
- else None,
738
- }
739
-
740
- # Wrap the message (first argument for publish, second for publish_to)
741
- message_idx = 1 if method_name == "publish_to" else 0
742
- if len(args) > message_idx:
743
- args_list = list(args)
744
- message = args_list[message_idx]
745
- wrapped_message = (
746
- SLIMInstrumentor._wrap_message_with_headers(
747
- None, message, headers
748
- )
749
- )
750
-
751
- if isinstance(wrapped_message, dict):
752
- message_to_send = json.dumps(wrapped_message).encode(
753
- "utf-8"
754
- )
755
- else:
756
- message_to_send = wrapped_message
757
-
758
- args_list[message_idx] = message_to_send
759
- args = tuple(args_list)
760
-
761
- return await original_method(self, *args, **kwargs)
762
-
763
- setattr(session_class, method_name, instrumented_session_publish)
764
-
765
- def _instrument_session_method_if_exists(self, slim_bindings, method_name):
766
- """Helper to instrument a session method if it exists"""
767
-
768
- # Look for session classes that might have this method
769
- for attr_name in dir(slim_bindings):
770
- attr = getattr(slim_bindings, attr_name)
771
- if hasattr(attr, method_name):
772
- original_method = getattr(attr, method_name)
773
-
774
- if callable(original_method):
775
- instrumented_method = self._create_session_method_wrapper(
776
- method_name, original_method
442
+ if args and len(args) > msg_idx and (traceparent or session_id):
443
+ headers = {
444
+ "session_id": session_id,
445
+ "traceparent": traceparent,
446
+ "slim_session_id": slim_session_id,
447
+ }
448
+ args_list = list(args)
449
+ wrapped_msg = _wrap_message_with_headers(
450
+ args_list[msg_idx], headers
777
451
  )
778
- setattr(attr, method_name, instrumented_method)
779
-
780
- def _create_session_method_wrapper(self, method_name, original_method):
781
- """Create an instrumented wrapper for session methods"""
782
-
783
- @functools.wraps(original_method)
784
- async def instrumented_session_method(self, *args, **kwargs):
785
- if _global_tracer:
786
- with _global_tracer.start_as_current_span(f"session.{method_name}"):
787
- traceparent = get_current_traceparent()
788
-
789
- # Handle message wrapping for publish methods
790
- if method_name in ["publish", "publish_to"] and args:
791
- # Thread-safe access to kv_store
792
- session_id = None
793
- if traceparent:
794
- with _kv_lock:
795
- session_id = kv_store.get(f"execution.{traceparent}")
796
- if session_id:
797
- kv_store.set(f"execution.{traceparent}", session_id)
798
-
799
- headers = {
800
- "session_id": session_id if session_id else None,
801
- "traceparent": traceparent,
802
- "slim_session_id": str(self.id)
803
- if hasattr(self, "id")
804
- else None,
805
- }
806
-
807
- # Set baggage context
808
- if traceparent and session_id:
809
- baggage.set_baggage(f"execution.{traceparent}", session_id)
810
-
811
- # Wrap the message (first argument for publish, second for publish_to)
812
- message_idx = 1 if method_name == "publish_to" else 0
813
- if len(args) > message_idx:
814
- args_list = list(args)
815
- message = args_list[message_idx]
816
- wrapped_message = SLIMInstrumentor._wrap_message_with_headers(
817
- None,
818
- message,
819
- headers, # Pass None for self since it's a static method
820
- )
821
-
822
- # Convert wrapped message back to bytes if needed
823
- if isinstance(wrapped_message, dict):
824
- message_to_send = json.dumps(wrapped_message).encode(
825
- "utf-8"
826
- )
827
- else:
828
- message_to_send = wrapped_message
829
-
830
- args_list[message_idx] = message_to_send
831
- args = tuple(args_list)
832
-
833
- return await original_method(self, *args, **kwargs)
834
- else:
835
- # Handle message wrapping even without tracing
836
- if method_name in ["publish", "publish_to"] and args:
837
- traceparent = get_current_traceparent()
838
- session_id = None
839
- if traceparent:
840
- with _kv_lock:
841
- session_id = kv_store.get(f"execution.{traceparent}")
842
-
843
- if traceparent or session_id:
844
- headers = {
845
- "session_id": session_id if session_id else None,
846
- "traceparent": traceparent,
847
- "slim_session_id": str(self.id)
848
- if hasattr(self, "id")
849
- else None,
850
- }
851
-
852
- # Wrap the message
853
- message_idx = 1 if method_name == "publish_to" else 0
854
- if len(args) > message_idx:
855
- args_list = list(args)
856
- message = args_list[message_idx]
857
- wrapped_message = (
858
- SLIMInstrumentor._wrap_message_with_headers(
859
- None, message, headers
860
- )
861
- )
862
-
863
- if isinstance(wrapped_message, dict):
864
- message_to_send = json.dumps(wrapped_message).encode(
865
- "utf-8"
866
- )
867
- else:
868
- message_to_send = wrapped_message
869
-
870
- args_list[message_idx] = message_to_send
871
- args = tuple(args_list)
452
+ args_list[msg_idx] = (
453
+ json.dumps(wrapped_msg).encode("utf-8")
454
+ if isinstance(wrapped_msg, dict)
455
+ else wrapped_msg
456
+ )
457
+ args = tuple(args_list)
872
458
 
873
- return await original_method(self, *args, **kwargs)
459
+ return await orig(self, *args, **kwargs)
874
460
 
875
- return instrumented_session_method
876
-
877
- @staticmethod
878
- def _wrap_message_with_headers(self, message, headers):
879
- """Helper method to wrap messages with headers consistently"""
880
- if isinstance(message, bytes):
881
- try:
882
- decoded_message = message.decode("utf-8")
883
- try:
884
- original_message = json.loads(decoded_message)
885
- if isinstance(original_message, dict):
886
- wrapped_message = original_message.copy()
887
- existing_headers = wrapped_message.get("headers", {})
888
- existing_headers.update(headers)
889
- wrapped_message["headers"] = existing_headers
890
- else:
891
- wrapped_message = {
892
- "headers": headers,
893
- "payload": original_message,
894
- }
895
- except json.JSONDecodeError:
896
- wrapped_message = {"headers": headers, "payload": decoded_message}
897
- except UnicodeDecodeError:
898
- # Fix type annotation issue by ensuring message is bytes
899
- encoded_message = (
900
- message if isinstance(message, bytes) else message.encode("utf-8")
901
- )
902
- wrapped_message = {
903
- "headers": headers,
904
- "payload": base64.b64encode(encoded_message).decode("utf-8"),
905
- }
906
- elif isinstance(message, str):
907
- try:
908
- original_message = json.loads(message)
909
- if isinstance(original_message, dict):
910
- wrapped_message = original_message.copy()
911
- existing_headers = wrapped_message.get("headers", {})
912
- existing_headers.update(headers)
913
- wrapped_message["headers"] = existing_headers
914
- else:
915
- wrapped_message = {"headers": headers, "payload": original_message}
916
- except json.JSONDecodeError:
917
- wrapped_message = {"headers": headers, "payload": message}
918
- elif isinstance(message, dict):
919
- wrapped_message = message.copy()
920
- existing_headers = wrapped_message.get("headers", {})
921
- existing_headers.update(headers)
922
- wrapped_message["headers"] = existing_headers
923
- else:
924
- wrapped_message = {"headers": headers, "payload": json.dumps(message)}
925
-
926
- return wrapped_message
461
+ setattr(session_class, method_name, wrapped)
927
462
 
928
463
  def _uninstrument(self, **kwargs):
929
464
  try:
930
465
  import slim_bindings
931
466
  except ImportError:
932
- raise ImportError(
933
- "No module named 'slim_bindings'. Please install it first."
467
+ return
468
+
469
+ def restore(obj, methods):
470
+ for m in methods:
471
+ if hasattr(obj, m):
472
+ orig = getattr(obj, m)
473
+ if hasattr(orig, "__wrapped__"):
474
+ setattr(obj, m, orig.__wrapped__)
475
+
476
+ if hasattr(slim_bindings, "Service"):
477
+ restore(slim_bindings.Service, ["connect_async", "run_server_async"])
478
+
479
+ if hasattr(slim_bindings, "App"):
480
+ restore(
481
+ slim_bindings.App,
482
+ [
483
+ "create_session_async",
484
+ "create_session_and_wait_async",
485
+ "subscribe_async",
486
+ "set_route_async",
487
+ "listen_for_session",
488
+ ],
934
489
  )
935
490
 
936
- # Restore the original methods
937
- methods_to_restore = [
938
- "publish",
939
- "publish_to",
940
- "request_reply", # v0.4.0-v0.5.x only
941
- "receive",
942
- "connect",
943
- "create_session",
944
- "invite",
945
- "set_route",
946
- "listen_for_session", # v0.6.0+
947
- ]
948
-
949
- for method_name in methods_to_restore:
950
- if hasattr(slim_bindings.Slim, method_name):
951
- original_method = getattr(slim_bindings.Slim, method_name)
952
- if hasattr(original_method, "__wrapped__"):
953
- setattr(
954
- slim_bindings.Slim, method_name, original_method.__wrapped__
955
- )
956
-
957
- # Also try to restore session-level methods (v0.6.0+)
958
- # This is best-effort since session classes may vary
959
- session_methods_to_restore = [
960
- "publish",
961
- "publish_to",
962
- "get_message",
963
- "invite",
964
- "remove",
965
- ]
966
-
967
- for attr_name in dir(slim_bindings):
968
- attr = getattr(slim_bindings, attr_name)
969
- for method_name in session_methods_to_restore:
970
- if hasattr(attr, method_name):
971
- original_method = getattr(attr, method_name)
972
- if hasattr(original_method, "__wrapped__"):
973
- setattr(attr, method_name, original_method.__wrapped__)
491
+ # Restore session methods
492
+ session_methods = ["publish_async", "publish_to_async", "get_message_async"]
493
+ for name in dir(slim_bindings):
494
+ cls = getattr(slim_bindings, name)
495
+ if isinstance(cls, type):
496
+ restore(cls, session_methods)