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