asap-protocol 0.5.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (59) hide show
  1. asap/__init__.py +1 -1
  2. asap/cli.py +137 -2
  3. asap/examples/README.md +81 -13
  4. asap/examples/auth_patterns.py +212 -0
  5. asap/examples/error_recovery.py +248 -0
  6. asap/examples/long_running.py +287 -0
  7. asap/examples/mcp_integration.py +240 -0
  8. asap/examples/multi_step_workflow.py +134 -0
  9. asap/examples/orchestration.py +293 -0
  10. asap/examples/rate_limiting.py +137 -0
  11. asap/examples/run_demo.py +0 -2
  12. asap/examples/secure_handler.py +84 -0
  13. asap/examples/state_migration.py +240 -0
  14. asap/examples/streaming_response.py +108 -0
  15. asap/examples/websocket_concept.py +129 -0
  16. asap/mcp/__init__.py +43 -0
  17. asap/mcp/client.py +224 -0
  18. asap/mcp/protocol.py +179 -0
  19. asap/mcp/server.py +333 -0
  20. asap/mcp/server_runner.py +40 -0
  21. asap/models/base.py +0 -3
  22. asap/models/constants.py +3 -1
  23. asap/models/entities.py +21 -6
  24. asap/models/envelope.py +7 -0
  25. asap/models/ids.py +8 -4
  26. asap/models/parts.py +33 -3
  27. asap/models/validators.py +16 -0
  28. asap/observability/__init__.py +6 -0
  29. asap/observability/dashboards/README.md +24 -0
  30. asap/observability/dashboards/asap-detailed.json +131 -0
  31. asap/observability/dashboards/asap-red.json +129 -0
  32. asap/observability/logging.py +81 -1
  33. asap/observability/metrics.py +15 -1
  34. asap/observability/trace_parser.py +238 -0
  35. asap/observability/trace_ui.py +218 -0
  36. asap/observability/tracing.py +293 -0
  37. asap/state/machine.py +15 -2
  38. asap/state/snapshot.py +0 -9
  39. asap/testing/__init__.py +31 -0
  40. asap/testing/assertions.py +108 -0
  41. asap/testing/fixtures.py +113 -0
  42. asap/testing/mocks.py +152 -0
  43. asap/transport/__init__.py +28 -0
  44. asap/transport/cache.py +180 -0
  45. asap/transport/circuit_breaker.py +9 -8
  46. asap/transport/client.py +418 -36
  47. asap/transport/compression.py +389 -0
  48. asap/transport/handlers.py +106 -53
  49. asap/transport/middleware.py +58 -34
  50. asap/transport/server.py +429 -139
  51. asap/transport/validators.py +0 -4
  52. asap/utils/sanitization.py +0 -5
  53. asap_protocol-1.0.0.dist-info/METADATA +264 -0
  54. asap_protocol-1.0.0.dist-info/RECORD +70 -0
  55. asap_protocol-0.5.0.dist-info/METADATA +0 -244
  56. asap_protocol-0.5.0.dist-info/RECORD +0 -41
  57. {asap_protocol-0.5.0.dist-info → asap_protocol-1.0.0.dist-info}/WHEEL +0 -0
  58. {asap_protocol-0.5.0.dist-info → asap_protocol-1.0.0.dist-info}/entry_points.txt +0 -0
  59. {asap_protocol-0.5.0.dist-info → asap_protocol-1.0.0.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,218 @@
1
+ """Optional Web UI for ASAP trace visualization.
2
+
3
+ Serves a simple FastAPI app to browse traces, search by trace_id, and visualize
4
+ request flow from pasted JSON log lines (ASAP_LOG_FORMAT=json).
5
+
6
+ Run with: uvicorn asap.observability.trace_ui:app --reload
7
+ Then open http://localhost:8000
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from typing import Any
13
+
14
+ from fastapi import FastAPI, HTTPException
15
+ from fastapi.responses import HTMLResponse
16
+ from pydantic import BaseModel
17
+
18
+ from asap.observability.trace_parser import (
19
+ TraceHop,
20
+ extract_trace_ids,
21
+ parse_trace_from_lines,
22
+ )
23
+
24
+ MAX_LOG_LINES_LENGTH = 2 * 1024 * 1024
25
+
26
+
27
+ def _parse_log_lines(raw: str, max_length: int = MAX_LOG_LINES_LENGTH) -> list[str]:
28
+ """Parse raw log lines into non-empty stripped lines; raise 413 if over max_length."""
29
+ if len(raw) > max_length:
30
+ raise HTTPException(
31
+ status_code=413,
32
+ detail=f"Log lines payload too large (max {max_length} bytes)",
33
+ )
34
+ return [s.strip() for s in raw.strip().splitlines() if s.strip()]
35
+
36
+
37
+ app = FastAPI(
38
+ title="ASAP Trace UI",
39
+ description="Browse and visualize ASAP request traces from JSON log lines.",
40
+ version="0.1.0",
41
+ )
42
+
43
+
44
+ class LogLinesBody(BaseModel):
45
+ """Request body with newline-separated log lines."""
46
+
47
+ log_lines: str
48
+
49
+
50
+ class VisualizeBody(BaseModel):
51
+ """Request body for trace visualization."""
52
+
53
+ log_lines: str
54
+ trace_id: str
55
+
56
+
57
+ def _hops_to_dict(hops: list[TraceHop]) -> list[dict[str, Any]]:
58
+ """Convert TraceHop list to JSON-serializable dicts."""
59
+ return [
60
+ {"sender": h.sender, "recipient": h.recipient, "duration_ms": h.duration_ms} for h in hops
61
+ ]
62
+
63
+
64
+ @app.get("/", response_class=HTMLResponse)
65
+ def index() -> HTMLResponse:
66
+ """Serve the trace browser UI (paste logs, list trace IDs, visualize)."""
67
+ html = _INDEX_HTML
68
+ return HTMLResponse(html)
69
+
70
+
71
+ @app.post("/api/traces/list")
72
+ def list_traces(body: LogLinesBody) -> dict[str, Any]:
73
+ """Extract unique trace IDs from pasted log lines."""
74
+ lines = _parse_log_lines(body.log_lines)
75
+ trace_ids = extract_trace_ids(lines)
76
+ return {"trace_ids": trace_ids}
77
+
78
+
79
+ @app.post("/api/traces/visualize")
80
+ def visualize_trace(body: VisualizeBody) -> dict[str, Any]:
81
+ """Parse logs and return hops + ASCII diagram for the given trace_id."""
82
+ lines = _parse_log_lines(body.log_lines)
83
+ hops, diagram = parse_trace_from_lines(lines, body.trace_id.strip())
84
+ if not diagram:
85
+ raise HTTPException(status_code=404, detail=f"No trace found for: {body.trace_id}")
86
+ return {
87
+ "trace_id": body.trace_id,
88
+ "hops": _hops_to_dict(hops),
89
+ "diagram": diagram,
90
+ }
91
+
92
+
93
+ _INDEX_HTML = """
94
+ <!DOCTYPE html>
95
+ <html lang="en">
96
+ <head>
97
+ <meta charset="utf-8">
98
+ <meta name="viewport" content="width=device-width, initial-scale=1">
99
+ <title>ASAP Trace UI</title>
100
+ <style>
101
+ * { box-sizing: border-box; }
102
+ body { font-family: system-ui, sans-serif; max-width: 900px; margin: 1rem auto; padding: 0 1rem; }
103
+ h1 { font-size: 1.25rem; }
104
+ label { display: block; margin-top: 0.75rem; font-weight: 600; }
105
+ textarea { width: 100%; min-height: 120px; font-family: monospace; font-size: 0.8rem; padding: 0.5rem; }
106
+ input[type="text"] { width: 100%; max-width: 320px; padding: 0.4rem; }
107
+ button { margin-top: 0.5rem; padding: 0.4rem 0.8rem; cursor: pointer; }
108
+ .result { margin-top: 1rem; padding: 0.75rem; background: #f5f5f5; border-radius: 4px; white-space: pre-wrap; font-family: monospace; font-size: 0.9rem; }
109
+ .error { background: #fee; color: #c00; }
110
+ ul { margin: 0.25rem 0; padding-left: 1.25rem; }
111
+ ul li { cursor: pointer; margin: 0.2rem 0; }
112
+ ul li:hover { text-decoration: underline; }
113
+ </style>
114
+ </head>
115
+ <body>
116
+ <h1>ASAP Trace UI</h1>
117
+ <p>Paste ASAP JSON log lines below (one event per line, <code>ASAP_LOG_FORMAT=json</code>), then list trace IDs or visualize one.</p>
118
+
119
+ <label for="logs">Log lines</label>
120
+ <textarea id="logs" placeholder='{"event":"asap.request.received","trace_id":"t1",...}
121
+ {"event":"asap.request.processed","trace_id":"t1","duration_ms":15,...}'></textarea>
122
+
123
+ <button id="btnList">List trace IDs</button>
124
+ <div id="listResult" class="result" style="display:none;"></div>
125
+
126
+ <label for="traceId">Trace ID (or click one above)</label>
127
+ <input type="text" id="traceId" placeholder="e.g. trace-abc">
128
+
129
+ <button id="btnViz">Visualize</button>
130
+ <div id="vizResult" class="result" style="display:none;"></div>
131
+
132
+ <script>
133
+ function escapeHtml(s) {
134
+ if (typeof s !== 'string') return '';
135
+ const m = { '&': '&amp;', '<': '&lt;', '>': '&gt;', '"': '&quot;', "'": '&#39;' };
136
+ return s.replace(/[&<>"']/g, function(c) { return m[c] || c; });
137
+ }
138
+ function errorMessage(e) {
139
+ return e instanceof Error ? e.message : String(e);
140
+ }
141
+ function showError(el, msg) {
142
+ el.textContent = 'Error: ' + msg;
143
+ el.classList.add('error');
144
+ el.style.display = 'block';
145
+ }
146
+ function renderTraceList(data, listResult, traceIdEl) {
147
+ if (data.trace_ids && data.trace_ids.length) {
148
+ const safe = data.trace_ids.map(function(t) {
149
+ return '<li data-trace="' + escapeHtml(t) + '">' + escapeHtml(t) + '</li>';
150
+ });
151
+ listResult.innerHTML = 'Trace IDs: <ul>' + safe.join('') + '</ul>';
152
+ listResult.querySelectorAll('li').forEach(function(li) {
153
+ li.onclick = function() { traceIdEl.value = li.dataset.trace || ''; };
154
+ });
155
+ } else {
156
+ listResult.textContent = 'No trace IDs found in log lines.';
157
+ }
158
+ listResult.classList.remove('error');
159
+ listResult.style.display = 'block';
160
+ }
161
+
162
+ const logsEl = document.getElementById('logs');
163
+ const traceIdEl = document.getElementById('traceId');
164
+ const listResult = document.getElementById('listResult');
165
+ const vizResult = document.getElementById('vizResult');
166
+
167
+ document.getElementById('btnList').onclick = async function() {
168
+ listResult.style.display = 'none';
169
+ const logLines = logsEl.value.trim();
170
+ if (!logLines) {
171
+ listResult.textContent = 'Paste log lines first.';
172
+ listResult.style.display = 'block';
173
+ return;
174
+ }
175
+ try {
176
+ const r = await fetch('/api/traces/list', {
177
+ method: 'POST',
178
+ headers: { 'Content-Type': 'application/json' },
179
+ body: JSON.stringify({ log_lines: logLines }),
180
+ });
181
+ const data = await r.json();
182
+ renderTraceList(data, listResult, traceIdEl);
183
+ } catch (e) {
184
+ showError(listResult, errorMessage(e));
185
+ }
186
+ };
187
+
188
+ document.getElementById('btnViz').onclick = async function() {
189
+ vizResult.style.display = 'none';
190
+ const logLines = logsEl.value.trim();
191
+ const traceId = traceIdEl.value.trim();
192
+ if (!logLines || !traceId) {
193
+ vizResult.textContent = 'Paste log lines and enter a trace ID.';
194
+ vizResult.style.display = 'block';
195
+ return;
196
+ }
197
+ try {
198
+ const r = await fetch('/api/traces/visualize', {
199
+ method: 'POST',
200
+ headers: { 'Content-Type': 'application/json' },
201
+ body: JSON.stringify({ log_lines: logLines, trace_id: traceId }),
202
+ });
203
+ if (!r.ok) {
204
+ const err = await r.json();
205
+ throw new Error(err.detail || r.statusText);
206
+ }
207
+ const data = await r.json();
208
+ vizResult.textContent = data.diagram || JSON.stringify(data, null, 2);
209
+ vizResult.classList.remove('error');
210
+ vizResult.style.display = 'block';
211
+ } catch (e) {
212
+ showError(vizResult, errorMessage(e));
213
+ }
214
+ };
215
+ </script>
216
+ </body>
217
+ </html>
218
+ """
@@ -0,0 +1,293 @@
1
+ """OpenTelemetry tracing integration for ASAP protocol.
2
+
3
+ This module provides distributed tracing with W3C Trace Context propagation.
4
+ When enabled, FastAPI and httpx are auto-instrumented; custom spans cover
5
+ handler execution and state transitions. Trace IDs can be carried in
6
+ envelope.trace_id and envelope.extensions for cross-service correlation.
7
+
8
+ Example:
9
+ >>> from asap.observability.tracing import configure_tracing, get_tracer
10
+ >>> configure_tracing(service_name="my-agent", app=app)
11
+ >>> tracer = get_tracer(__name__)
12
+ >>> with tracer.start_as_current_span("my.operation"):
13
+ ... ...
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import os
19
+ from contextlib import AbstractContextManager
20
+ from typing import TYPE_CHECKING, Any
21
+
22
+ from opentelemetry import context, trace
23
+ from opentelemetry.sdk.resources import Resource
24
+ from opentelemetry.sdk.trace import TracerProvider
25
+ from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
26
+ from opentelemetry.sdk.trace.export import SimpleSpanProcessor
27
+
28
+ if TYPE_CHECKING:
29
+ from fastapi import FastAPI
30
+
31
+ from asap.models.envelope import Envelope
32
+ from asap.observability.logging import get_logger
33
+
34
+ # Environment variables for zero-config (OpenTelemetry convention)
35
+ _ENV_OTEL_SERVICE_NAME = "OTEL_SERVICE_NAME"
36
+ _ENV_OTEL_TRACES_EXPORTER = "OTEL_TRACES_EXPORTER"
37
+ _ENV_OTEL_EXPORTER_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT"
38
+
39
+ # Extension keys for W3C trace context in envelope
40
+ EXTENSION_TRACE_ID = "trace_id"
41
+ EXTENSION_SPAN_ID = "span_id"
42
+
43
+ logger = get_logger(__name__)
44
+
45
+ _tracer_provider: TracerProvider | None = None
46
+ _tracer: trace.Tracer | None = None
47
+
48
+
49
+ def configure_tracing(
50
+ service_name: str | None = None,
51
+ app: FastAPI | None = None,
52
+ ) -> None:
53
+ """Configure OpenTelemetry tracing and optionally instrument FastAPI and httpx.
54
+
55
+ Uses environment variables for zero-config:
56
+ - OTEL_SERVICE_NAME: service name (default: "asap-server")
57
+ - OTEL_TRACES_EXPORTER: "none" | "otlp" | "console" (default: "none")
58
+ - OTEL_EXPORTER_OTLP_ENDPOINT: OTLP endpoint (e.g. http://localhost:4317)
59
+
60
+ When OTEL_TRACES_EXPORTER is "none" or unset, tracing is configured but
61
+ no spans are exported (useful for dev without Jaeger). Set to "otlp" and
62
+ OTEL_EXPORTER_OTLP_ENDPOINT for Jaeger/collector.
63
+
64
+ Args:
65
+ service_name: Override for OTEL_SERVICE_NAME.
66
+ app: If provided, FastAPI and httpx are instrumented.
67
+ """
68
+ global _tracer_provider, _tracer
69
+
70
+ name = service_name or os.environ.get(_ENV_OTEL_SERVICE_NAME) or "asap-server"
71
+ resource = Resource.create({"service.name": name})
72
+ _tracer_provider = TracerProvider(resource=resource)
73
+ trace.set_tracer_provider(_tracer_provider)
74
+
75
+ exporter_name = os.environ.get(_ENV_OTEL_TRACES_EXPORTER, "none").strip().lower()
76
+ if exporter_name == "none":
77
+ _tracer = trace.get_tracer("asap.protocol", "1.0.0", schema_url=None)
78
+ if app is not None:
79
+ _instrument_app(app)
80
+ return
81
+
82
+ if exporter_name == "otlp":
83
+ _add_otlp_processor()
84
+ elif exporter_name == "console":
85
+ _add_console_processor()
86
+
87
+ _tracer = trace.get_tracer("asap.protocol", "1.0.0", schema_url=None)
88
+ if app is not None:
89
+ _instrument_app(app)
90
+
91
+
92
+ def _add_otlp_processor() -> None:
93
+ """Add OTLP span processor if endpoint is set."""
94
+ global _tracer_provider
95
+ if _tracer_provider is None:
96
+ return
97
+ endpoint = os.environ.get(_ENV_OTEL_EXPORTER_OTLP_ENDPOINT)
98
+ if not endpoint:
99
+ return
100
+ try:
101
+ from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
102
+
103
+ exporter: SpanExporter = OTLPSpanExporter(endpoint=endpoint, insecure=True)
104
+ _tracer_provider.add_span_processor(BatchSpanProcessor(exporter))
105
+ except ImportError as e:
106
+ logger.debug("OTLP gRPC exporter not available: %s", e)
107
+ try:
108
+ from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
109
+ OTLPSpanExporter as OTLPSpanExporterHttp,
110
+ )
111
+
112
+ exporter = OTLPSpanExporterHttp(endpoint=endpoint)
113
+ _tracer_provider.add_span_processor(BatchSpanProcessor(exporter))
114
+ except ImportError as e2:
115
+ logger.debug("OTLP HTTP exporter not available: %s", e2)
116
+
117
+
118
+ def _add_console_processor() -> None:
119
+ """Add console span processor for debugging."""
120
+ global _tracer_provider
121
+ if _tracer_provider is None:
122
+ return
123
+ try:
124
+ from opentelemetry.sdk.trace.export import ConsoleSpanExporter
125
+
126
+ _tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
127
+ except ImportError as e:
128
+ logger.debug("Console span exporter not available: %s", e)
129
+
130
+
131
+ def _instrument_app(app: FastAPI) -> None:
132
+ """Instrument FastAPI and httpx for automatic spans."""
133
+ try:
134
+ from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
135
+
136
+ FastAPIInstrumentor.instrument_app(app)
137
+ except Exception as e:
138
+ logger.debug("Failed to instrument FastAPI: %s", e)
139
+ try:
140
+ from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
141
+
142
+ HTTPXClientInstrumentor().instrument()
143
+ except Exception as e:
144
+ logger.debug("Failed to instrument httpx: %s", e)
145
+
146
+
147
+ def reset_tracing() -> None:
148
+ """Reset global tracer state (for test teardown).
149
+
150
+ Clears the module-level tracer provider and tracer so that subsequent
151
+ tests or configure_tracing() calls start from a clean state.
152
+ """
153
+ global _tracer_provider, _tracer
154
+ _tracer_provider = None
155
+ _tracer = None
156
+
157
+
158
+ def get_tracer(name: str | None = None) -> trace.Tracer:
159
+ """Return the ASAP protocol tracer for custom spans.
160
+
161
+ Args:
162
+ name: Optional logger/module name for span attribution.
163
+
164
+ Returns:
165
+ OpenTelemetry Tracer instance.
166
+ """
167
+ if _tracer is None:
168
+ trace.set_tracer_provider(TracerProvider())
169
+ return trace.get_tracer("asap.protocol", "1.0.0", schema_url=None)
170
+ return _tracer
171
+
172
+
173
+ def inject_envelope_trace_context(envelope: Envelope) -> Envelope:
174
+ """Inject current span trace_id and span_id into envelope (W3C propagation).
175
+
176
+ Sets envelope.trace_id and envelope.extensions[trace_id, span_id] from the
177
+ current OpenTelemetry context so downstream services can continue the trace.
178
+
179
+ Args:
180
+ envelope: Response envelope to annotate.
181
+
182
+ Returns:
183
+ New envelope with trace_id and extensions updated (immutable).
184
+ """
185
+ span = trace.get_current_span()
186
+ if not span.is_recording():
187
+ return envelope
188
+
189
+ ctx = span.get_span_context()
190
+ trace_id_hex = format(ctx.trace_id, "032x")
191
+ span_id_hex = format(ctx.span_id, "016x")
192
+
193
+ extensions = dict(envelope.extensions or {})
194
+ extensions[EXTENSION_TRACE_ID] = trace_id_hex
195
+ extensions[EXTENSION_SPAN_ID] = span_id_hex
196
+
197
+ # Preserve existing trace_id for correlation (e.g. from request); set only if missing
198
+ new_trace_id = envelope.trace_id if envelope.trace_id else trace_id_hex
199
+ return envelope.model_copy(update={"trace_id": new_trace_id, "extensions": extensions})
200
+
201
+
202
+ def extract_and_activate_envelope_trace_context(envelope: Envelope) -> Any | None:
203
+ """Extract trace context from envelope and set as current context (W3C).
204
+
205
+ If envelope has trace_id (and optionally span_id in extensions), creates
206
+ a non-recording span context so new spans become children of the incoming
207
+ trace. Caller should call context.detach(token) when request ends.
208
+
209
+ Args:
210
+ envelope: Incoming envelope carrying trace_id / extensions.
211
+
212
+ Returns:
213
+ Context token to pass to context.detach(), or None if no context set.
214
+ """
215
+ trace_id_str = envelope.trace_id
216
+ if not trace_id_str or len(trace_id_str) != 32:
217
+ return None
218
+
219
+ extensions = envelope.extensions or {}
220
+ span_id_str = extensions.get(EXTENSION_SPAN_ID)
221
+ if not span_id_str or len(span_id_str) != 16:
222
+ return None
223
+
224
+ try:
225
+ trace_id = int(trace_id_str, 16)
226
+ span_id = int(span_id_str, 16)
227
+ except ValueError:
228
+ return None
229
+
230
+ from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
231
+
232
+ span_context = SpanContext(
233
+ trace_id=trace_id,
234
+ span_id=span_id,
235
+ is_remote=True,
236
+ trace_flags=TraceFlags(0x01),
237
+ )
238
+ ctx = trace.set_span_in_context(NonRecordingSpan(span_context))
239
+ return context.attach(ctx)
240
+
241
+
242
+ def handler_span_context(
243
+ payload_type: str,
244
+ agent_urn: str,
245
+ envelope_id: str | None,
246
+ ) -> AbstractContextManager[trace.Span]:
247
+ """Start a current span for handler execution (use as context manager).
248
+
249
+ Attributes: asap.payload_type, asap.agent.urn, asap.envelope.id.
250
+
251
+ Args:
252
+ payload_type: Envelope payload type.
253
+ agent_urn: Manifest/agent URN.
254
+ envelope_id: Envelope id.
255
+
256
+ Returns:
257
+ Span context manager (use with "with").
258
+ """
259
+ tracer = get_tracer(__name__)
260
+ attrs: dict[str, str] = {
261
+ "asap.payload_type": payload_type,
262
+ "asap.agent.urn": agent_urn,
263
+ }
264
+ if envelope_id:
265
+ attrs["asap.envelope.id"] = envelope_id
266
+ return tracer.start_as_current_span("asap.handler.execute", attributes=attrs)
267
+
268
+
269
+ def state_transition_span_context(
270
+ from_status: str,
271
+ to_status: str,
272
+ task_id: str | None = None,
273
+ ) -> AbstractContextManager[trace.Span]:
274
+ """Start a current span for a state machine transition (use as context manager).
275
+
276
+ Attributes: asap.state.from, asap.state.to, asap.task.id.
277
+
278
+ Args:
279
+ from_status: Previous status.
280
+ to_status: New status.
281
+ task_id: Optional task id.
282
+
283
+ Returns:
284
+ Span context manager (use with "with").
285
+ """
286
+ tracer = get_tracer(__name__)
287
+ attrs: dict[str, str] = {
288
+ "asap.state.from": from_status,
289
+ "asap.state.to": to_status,
290
+ }
291
+ if task_id:
292
+ attrs["asap.task.id"] = task_id
293
+ return tracer.start_as_current_span("asap.state.transition", attributes=attrs)
asap/state/machine.py CHANGED
@@ -14,7 +14,10 @@ from datetime import datetime, timezone
14
14
  from asap.errors import InvalidTransitionError
15
15
  from asap.models.entities import Task
16
16
  from asap.models.enums import TaskStatus
17
+ from asap.observability import get_metrics
18
+ from asap.observability.tracing import state_transition_span_context
17
19
 
20
+ __all__ = ["TaskStatus", "can_transition", "transition", "VALID_TRANSITIONS"]
18
21
 
19
22
  # Valid state transitions mapping
20
23
  VALID_TRANSITIONS: dict[TaskStatus, set[TaskStatus]] = {
@@ -82,5 +85,15 @@ def transition(task: Task, new_status: TaskStatus) -> Task:
82
85
  from_state=task.status.value, to_state=new_status.value, details={"task_id": task.id}
83
86
  )
84
87
 
85
- # Create new task instance with updated status and timestamp (immutable approach)
86
- return task.model_copy(update={"status": new_status, "updated_at": datetime.now(timezone.utc)})
88
+ with state_transition_span_context(
89
+ from_status=task.status.value,
90
+ to_status=new_status.value,
91
+ task_id=task.id,
92
+ ):
93
+ get_metrics().increment_counter(
94
+ "asap_state_transitions_total",
95
+ {"from_status": task.status.value, "to_status": new_status.value},
96
+ )
97
+ return task.model_copy(
98
+ update={"status": new_status, "updated_at": datetime.now(timezone.utc)}
99
+ )
asap/state/snapshot.py CHANGED
@@ -154,14 +154,10 @@ class InMemorySnapshotStore:
154
154
  with self._lock:
155
155
  task_id = snapshot.task_id
156
156
 
157
- # Initialize storage for this task if needed
158
157
  if task_id not in self._snapshots:
159
158
  self._snapshots[task_id] = {}
160
159
 
161
- # Store the snapshot
162
160
  self._snapshots[task_id][snapshot.version] = snapshot
163
-
164
- # Update latest version
165
161
  self._latest_versions[task_id] = max(
166
162
  self._latest_versions.get(task_id, 0), snapshot.version
167
163
  )
@@ -186,13 +182,11 @@ class InMemorySnapshotStore:
186
182
  return None
187
183
 
188
184
  if version is None:
189
- # Return latest version
190
185
  latest_version = self._latest_versions.get(task_id)
191
186
  if latest_version is None:
192
187
  return None
193
188
  return self._snapshots[task_id].get(latest_version)
194
189
 
195
- # Return specific version
196
190
  return self._snapshots[task_id].get(version)
197
191
 
198
192
  def list_versions(self, task_id: TaskID) -> list[int]:
@@ -243,18 +237,15 @@ class InMemorySnapshotStore:
243
237
  return True
244
238
  return False
245
239
 
246
- # Delete specific version
247
240
  if version in self._snapshots[task_id]:
248
241
  del self._snapshots[task_id][version]
249
242
 
250
- # Update latest version if needed
251
243
  if self._latest_versions.get(task_id) == version:
252
244
  if self._snapshots[task_id]:
253
245
  self._latest_versions[task_id] = max(self._snapshots[task_id].keys())
254
246
  else:
255
247
  del self._latest_versions[task_id]
256
248
 
257
- # Clean up empty task dict
258
249
  if not self._snapshots[task_id]:
259
250
  del self._snapshots[task_id]
260
251
  if task_id in self._latest_versions:
@@ -0,0 +1,31 @@
1
+ """ASAP testing utilities for easier test authoring.
2
+
3
+ This package provides pytest fixtures, mock agents, and custom assertions
4
+ to reduce boilerplate when testing ASAP protocol integrations.
5
+
6
+ Modules:
7
+ fixtures: Pytest fixtures (mock_agent, mock_client, mock_snapshot_store)
8
+ and context managers (test_agent, test_client).
9
+ mocks: MockAgent for configurable mock agents with pre-set responses
10
+ and request recording.
11
+ assertions: Custom assertions (assert_envelope_valid, assert_task_completed,
12
+ assert_response_correlates).
13
+
14
+ Example:
15
+ >>> from asap.testing import MockAgent, assert_envelope_valid
16
+ >>> from asap.testing.fixtures import mock_agent, test_client
17
+ """
18
+
19
+ from asap.testing.assertions import (
20
+ assert_envelope_valid,
21
+ assert_response_correlates,
22
+ assert_task_completed,
23
+ )
24
+ from asap.testing.mocks import MockAgent
25
+
26
+ __all__ = [
27
+ "MockAgent",
28
+ "assert_envelope_valid",
29
+ "assert_response_correlates",
30
+ "assert_task_completed",
31
+ ]