Skip to content

Commit 10d22a2

Browse files
NiteshDhanpalclaude
NiteshDhanpal
and
claude
authored
feat(tracing): skip Agentex span-start write by default (end-only ingest) (#438)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 2d63eef commit 10d22a2

2 files changed

Lines changed: 250 additions & 25 deletions

File tree

‎src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py‎

Lines changed: 96 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
1+
import os
12
import asyncio
23
import weakref
34
from typing import TYPE_CHECKING, Any, Dict, override
45

56
from agentex import Agentex
67
from agentex.types.span import Span
78
from agentex.lib.types.tracing import AgentexTracingProcessorConfig
9+
from agentex.lib.utils.logging import make_logger
810
from agentex.lib.adk.utils._modules.client import create_async_agentex_client
911
from agentex.lib.core.tracing.processors.tracing_processor_interface import (
1012
SyncTracingProcessor,
@@ -14,28 +16,88 @@
1416
if TYPE_CHECKING:
1517
from agentex import AsyncAgentex
1618

19+
logger = make_logger(__name__)
20+
21+
22+
# NOTE: This is the Agentex-backend toggle (writes to the agentex `spans`
23+
# table via the Agentex API). It is intentionally SEPARATE from the SGP/EGP
24+
# processor's ``AGENTEX_TRACING_SKIP_SPAN_START`` so the two backends can be
25+
# controlled independently.
26+
_SKIP_SPAN_START_ENV = "AGENTEX_TRACING_SKIP_AGENTEX_SPAN_START"
27+
28+
29+
def _skip_span_start_enabled() -> bool:
30+
"""Whether to skip the Agentex span-start write and persist each span only on end.
31+
32+
The Agentex processor otherwise writes every span twice: a ``spans.create``
33+
on start (no ``end_time``/``output`` yet) and a ``spans.update`` on end.
34+
The start row is overwritten by the end write moments later, so persisting
35+
it doubles the per-span HTTP/DB write volume against the Agentex control
36+
plane — the load that timed out span-start activities and pressured the
37+
Agentex Postgres connection pool under load.
38+
39+
When enabled (the default), the start write is skipped and the END write
40+
becomes a single ``spans.create`` carrying the complete span — one INSERT
41+
per span instead of an INSERT + UPDATE. (A plain ``spans.update`` on end
42+
would 404 because the row was never created.)
43+
44+
Default ON. Set ``AGENTEX_TRACING_SKIP_AGENTEX_SPAN_START`` to
45+
``0``/``false``/``no``/``off`` to restore the start write — e.g. if you
46+
need in-flight spans visible before they complete, or spans that never end
47+
(process crash) to still be persisted.
48+
"""
49+
raw = os.environ.get(_SKIP_SPAN_START_ENV, "1").strip().lower()
50+
return raw not in ("0", "false", "no", "off")
51+
52+
53+
def _create_kwargs(span: Span) -> Dict[str, Any]:
54+
"""Full-span kwargs for ``spans.create`` — used on start (skip disabled) and
55+
on end (skip enabled, single-INSERT path)."""
56+
return {
57+
"name": span.name,
58+
"start_time": span.start_time,
59+
"end_time": span.end_time,
60+
"id": span.id,
61+
"trace_id": span.trace_id,
62+
"parent_id": span.parent_id,
63+
"input": span.input,
64+
"output": span.output,
65+
"data": span.data,
66+
"task_id": span.task_id,
67+
}
68+
1769

1870
class AgentexSyncTracingProcessor(SyncTracingProcessor):
1971
def __init__(self, config: AgentexTracingProcessorConfig): # noqa: ARG002
2072
self.client = Agentex()
73+
# Capture the skip decision once at init: both halves of a span's
74+
# lifecycle MUST agree, otherwise a start-skip + end-update lands on a
75+
# non-existent row (404) — or the reverse double-creates. Re-reading the
76+
# env per event would let a mid-span toggle (tests, config reload) split
77+
# the decision. Deploy-time flag, so a single read is correct.
78+
self._skip_span_start = _skip_span_start_enabled()
79+
logger.info(
80+
"Agentex tracing span-start write %s (%s)",
81+
"disabled — end-only ingest" if self._skip_span_start else "enabled",
82+
_SKIP_SPAN_START_ENV,
83+
)
2184

2285
@override
2386
def on_span_start(self, span: Span) -> None:
24-
self.client.spans.create(
25-
name=span.name,
26-
start_time=span.start_time,
27-
end_time=span.end_time,
28-
trace_id=span.trace_id,
29-
id=span.id,
30-
data=span.data,
31-
input=span.input,
32-
output=span.output,
33-
parent_id=span.parent_id,
34-
task_id=span.task_id,
35-
)
87+
# End-only ingest: by default the start write is skipped (see
88+
# _skip_span_start_enabled) so each span is persisted once, on end.
89+
if self._skip_span_start:
90+
return
91+
self.client.spans.create(**_create_kwargs(span))
3692

3793
@override
3894
def on_span_end(self, span: Span) -> None:
95+
# End-only ingest: the start create was skipped, so persist the complete
96+
# span as a single INSERT here (a bare spans.update would 404 — no row).
97+
if self._skip_span_start:
98+
self.client.spans.create(**_create_kwargs(span))
99+
return
100+
39101
update: Dict[str, Any] = {}
40102
if span.trace_id:
41103
update["trace_id"] = span.trace_id
@@ -82,6 +144,17 @@ def __init__(self, config: AgentexTracingProcessorConfig): # noqa: ARG002
82144
self._clients_by_loop: weakref.WeakKeyDictionary[
83145
asyncio.AbstractEventLoop, "AsyncAgentex"
84146
] = weakref.WeakKeyDictionary()
147+
# Capture the skip decision once at init: both halves of a span's
148+
# lifecycle MUST agree, otherwise a start-skip + end-update lands on a
149+
# non-existent row (404) — or the reverse double-creates. Re-reading the
150+
# env per event would let a mid-span toggle (tests, config reload) split
151+
# the decision. Deploy-time flag, so a single read is correct.
152+
self._skip_span_start = _skip_span_start_enabled()
153+
logger.info(
154+
"Agentex tracing span-start write %s (%s)",
155+
"disabled — end-only ingest" if self._skip_span_start else "enabled",
156+
_SKIP_SPAN_START_ENV,
157+
)
85158

86159
def _build_client(self) -> "AsyncAgentex":
87160
import httpx
@@ -111,21 +184,20 @@ def client(self) -> "AsyncAgentex":
111184
# https://linear.app/scale-epd/issue/AGX1-199/add-agentex-batch-endpoint-for-traces
112185
@override
113186
async def on_span_start(self, span: Span) -> None:
114-
await self.client.spans.create(
115-
name=span.name,
116-
start_time=span.start_time,
117-
end_time=span.end_time,
118-
id=span.id,
119-
trace_id=span.trace_id,
120-
parent_id=span.parent_id,
121-
input=span.input,
122-
output=span.output,
123-
data=span.data,
124-
task_id=span.task_id,
125-
)
187+
# End-only ingest: by default the start write is skipped (see
188+
# _skip_span_start_enabled) so each span is persisted once, on end.
189+
if self._skip_span_start:
190+
return
191+
await self.client.spans.create(**_create_kwargs(span))
126192

127193
@override
128194
async def on_span_end(self, span: Span) -> None:
195+
# End-only ingest: the start create was skipped, so persist the complete
196+
# span as a single INSERT here (a bare spans.update would 404 — no row).
197+
if self._skip_span_start:
198+
await self.client.spans.create(**_create_kwargs(span))
199+
return
200+
129201
update: Dict[str, Any] = {}
130202
if span.trace_id:
131203
update["trace_id"] = span.trace_id

‎tests/lib/core/tracing/processors/test_agentex_tracing_processor.py‎

Lines changed: 154 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22

33
import asyncio
44
import weakref
5-
from unittest.mock import MagicMock, patch
5+
from datetime import datetime, timezone
6+
from unittest.mock import AsyncMock, MagicMock, patch
67

78
import pytest
89

@@ -24,11 +25,163 @@
2425
MODULE = "agentex.lib.core.tracing.processors.agentex_tracing_processor"
2526

2627

28+
SKIP_ENV = "AGENTEX_TRACING_SKIP_AGENTEX_SPAN_START"
29+
30+
2731
def _make_config() -> MagicMock:
2832
"""Empty config — AgentexTracingProcessorConfig is unused by __init__."""
2933
return MagicMock()
3034

3135

36+
def _make_span():
37+
from agentex.types.span import Span
38+
39+
now = datetime.now(timezone.utc)
40+
return Span(
41+
id="span-1",
42+
trace_id="trace-1",
43+
name="test-span",
44+
start_time=now,
45+
end_time=now,
46+
input={"in": 1},
47+
output={"out": 2},
48+
)
49+
50+
51+
class TestAgentexSyncSkipSpanStart:
52+
"""The Agentex backend writes create-on-start + update-on-end by default.
53+
End-only ingest (default) skips the start write and makes the END a single
54+
create — verify the start is a no-op and end does an INSERT, not an UPDATE.
55+
"""
56+
57+
def test_start_skipped_and_end_creates_by_default(self, monkeypatch):
58+
monkeypatch.delenv(SKIP_ENV, raising=False) # default ON
59+
with patch(f"{MODULE}.Agentex") as MockAgentex:
60+
from agentex.lib.core.tracing.processors.agentex_tracing_processor import (
61+
AgentexSyncTracingProcessor,
62+
)
63+
64+
processor = AgentexSyncTracingProcessor(_make_config())
65+
client = MockAgentex.return_value
66+
span = _make_span()
67+
68+
processor.on_span_start(span)
69+
client.spans.create.assert_not_called() # start skipped
70+
client.spans.update.assert_not_called()
71+
72+
processor.on_span_end(span)
73+
client.spans.create.assert_called_once() # single INSERT on end
74+
client.spans.update.assert_not_called() # never a 404-prone UPDATE
75+
76+
def test_start_creates_and_end_updates_when_skip_disabled(self, monkeypatch):
77+
monkeypatch.setenv(SKIP_ENV, "0")
78+
with patch(f"{MODULE}.Agentex") as MockAgentex:
79+
from agentex.lib.core.tracing.processors.agentex_tracing_processor import (
80+
AgentexSyncTracingProcessor,
81+
)
82+
83+
processor = AgentexSyncTracingProcessor(_make_config())
84+
client = MockAgentex.return_value
85+
span = _make_span()
86+
87+
processor.on_span_start(span)
88+
client.spans.create.assert_called_once() # start write restored
89+
90+
processor.on_span_end(span)
91+
client.spans.update.assert_called_once() # end is the UPDATE
92+
93+
def test_skip_decision_captured_at_init_not_per_call(self, monkeypatch):
94+
"""The two halves of a span MUST use the same skip decision. A flag
95+
toggled after construction must not split it (start-skip + end-update
96+
would 404). The decision is captured once at init.
97+
"""
98+
monkeypatch.delenv(SKIP_ENV, raising=False) # construct with skip ON
99+
with patch(f"{MODULE}.Agentex") as MockAgentex:
100+
from agentex.lib.core.tracing.processors.agentex_tracing_processor import (
101+
AgentexSyncTracingProcessor,
102+
)
103+
104+
processor = AgentexSyncTracingProcessor(_make_config())
105+
client = MockAgentex.return_value
106+
span = _make_span()
107+
108+
processor.on_span_start(span) # skipped (cached ON)
109+
monkeypatch.setenv(SKIP_ENV, "0") # toggle mid-span — must be ignored
110+
processor.on_span_end(span)
111+
112+
client.spans.create.assert_called_once() # still end-only INSERT
113+
client.spans.update.assert_not_called() # NOT a 404-prone UPDATE
114+
115+
116+
class TestAgentexAsyncSkipSpanStart:
117+
async def test_start_skipped_and_end_creates_by_default(self, monkeypatch):
118+
monkeypatch.delenv(SKIP_ENV, raising=False) # default ON
119+
with patch(f"{MODULE}.create_async_agentex_client") as mock_factory:
120+
client = MagicMock()
121+
client.spans.create = AsyncMock()
122+
client.spans.update = AsyncMock()
123+
mock_factory.return_value = client
124+
125+
from agentex.lib.core.tracing.processors.agentex_tracing_processor import (
126+
AgentexAsyncTracingProcessor,
127+
)
128+
129+
processor = AgentexAsyncTracingProcessor(_make_config())
130+
span = _make_span()
131+
132+
await processor.on_span_start(span)
133+
client.spans.create.assert_not_called() # start skipped
134+
client.spans.update.assert_not_called()
135+
136+
await processor.on_span_end(span)
137+
client.spans.create.assert_awaited_once() # single INSERT on end
138+
client.spans.update.assert_not_called()
139+
140+
async def test_start_creates_and_end_updates_when_skip_disabled(self, monkeypatch):
141+
monkeypatch.setenv(SKIP_ENV, "0")
142+
with patch(f"{MODULE}.create_async_agentex_client") as mock_factory:
143+
client = MagicMock()
144+
client.spans.create = AsyncMock()
145+
client.spans.update = AsyncMock()
146+
mock_factory.return_value = client
147+
148+
from agentex.lib.core.tracing.processors.agentex_tracing_processor import (
149+
AgentexAsyncTracingProcessor,
150+
)
151+
152+
processor = AgentexAsyncTracingProcessor(_make_config())
153+
span = _make_span()
154+
155+
await processor.on_span_start(span)
156+
client.spans.create.assert_awaited_once() # start write restored
157+
158+
await processor.on_span_end(span)
159+
client.spans.update.assert_awaited_once() # end is the UPDATE
160+
161+
async def test_skip_decision_captured_at_init_not_per_call(self, monkeypatch):
162+
"""A flag toggled after construction must not split a span's lifecycle."""
163+
monkeypatch.delenv(SKIP_ENV, raising=False) # construct with skip ON
164+
with patch(f"{MODULE}.create_async_agentex_client") as mock_factory:
165+
client = MagicMock()
166+
client.spans.create = AsyncMock()
167+
client.spans.update = AsyncMock()
168+
mock_factory.return_value = client
169+
170+
from agentex.lib.core.tracing.processors.agentex_tracing_processor import (
171+
AgentexAsyncTracingProcessor,
172+
)
173+
174+
processor = AgentexAsyncTracingProcessor(_make_config())
175+
span = _make_span()
176+
177+
await processor.on_span_start(span) # skipped (cached ON)
178+
monkeypatch.setenv(SKIP_ENV, "0") # toggle mid-span — must be ignored
179+
await processor.on_span_end(span)
180+
181+
client.spans.create.assert_awaited_once() # still end-only INSERT
182+
client.spans.update.assert_not_called() # NOT a 404-prone UPDATE
183+
184+
32185
class TestAgentexAsyncTracingProcessor:
33186
"""Coverage for the per-event-loop client cache. The SGP processor has
34187
matching tests; mirror them here so a regression in the Agentex side

0 commit comments

Comments
 (0)