Merge pull request #31 from agentevals-dev/fix/live-session-groupping

Fix incorrect span groupping on UI

Author: krisztianfekete

src/agentevals/streaming/session.py

@@ -20,6 +20,7 @@ class TraceSession:
     logs: list[dict] = field(default_factory=list)
     started_at: datetime = field(default_factory=lambda: datetime.now(UTC))
     is_complete: bool = False
+    completed_at: datetime | None = None
     metadata: dict = field(default_factory=dict)
     source: str = "websocket"
     has_root_span: bool = False

src/agentevals/streaming/ws_server.py

@@ -233,6 +233,11 @@ async def get_or_create_otlp_session(
                 active.trace_ids.add(trace_id)
                 return active
 
+        existing = self.find_session_by_trace_id(trace_id)
+        if existing and existing.is_complete:
+            self._reopen_session(existing, trace_id, session_name)
+            return existing
+
         session_id = session_name
         if session_id in self.sessions:
             counter = 2
@@ -323,6 +328,29 @@ def schedule_log_reextraction(self, session_id: str) -> None:
             self._delayed_reextract(session_id, self.reextraction_delay_seconds)
         )
 
+    def _reopen_session(
+        self, session: TraceSession, trace_id: str, session_name: str
+    ) -> None:
+        """Reopen a completed session when a trace_id already in the session
+        receives more spans after completion (split-batch scenario).
+
+        The OTLP BatchSpanProcessor may flush one turn's spans across the
+        completion boundary: some child spans arrive before the grace period
+        fires, and the root span (plus remaining children) arrives after.
+        Because the trace_id was already registered in the session, we know
+        these late spans belong here rather than to a new agent run.
+        """
+        session.is_complete = False
+        session.completed_at = None
+        session.trace_ids.add(trace_id)
+        self._active_session_for_name[session_name] = session.session_id
+        self.incremental_extractors[session.session_id] = IncrementalInvocationExtractor()
+        self.reset_idle_timer(session.session_id)
+        logger.info(
+            "Reopened session %s for trace %s (%d spans so far)",
+            session.session_id, trace_id, len(session.spans),
+        )
+
     async def _delayed_complete(self, session_id: str, delay: float) -> None:
         await asyncio.sleep(delay)
         await self._complete_otlp_session(session_id)
@@ -376,6 +404,7 @@ async def _complete_otlp_session(self, session_id: str) -> None:
             return
 
         session.is_complete = True
+        session.completed_at = datetime.now(UTC)
 
         for name, sid in list(self._active_session_for_name.items()):
             if sid == session_id:

tests/integration/test_live_agents.py

@@ -187,3 +187,107 @@ def test_session_visible_via_api(self, live_servers):
         assert resp.status_code == 200
         session_ids = [s["sessionId"] for s in resp.json()["data"]]
         assert session_name in session_ids
+
+
+@_skip_no_openai
+class TestAgentRerun:
+    """Verify that re-running an agent with the same session_name creates
+    separate sessions, not merging them into one.
+
+    Each subprocess is a new OS process with new trace_ids. The OTLP
+    receiver must recognize these as distinct runs and assign unique
+    session IDs (session_name, session_name-2, etc.)."""
+
+    def test_strands_rerun_creates_separate_sessions(self, live_servers):
+        """Run the Strands agent twice with the same session_name.
+        Each run must produce its own session."""
+        main_port, otlp_port, mgr = live_servers
+        session_name = "e2e-strands-rerun"
+        strands_env = {
+            "OTEL_SEMCONV_STABILITY_OPT_IN": "gen_ai_latest_experimental",
+        }
+
+        result1 = _run_agent(
+            "examples/zero-code-examples/strands/run.py",
+            otlp_port,
+            session_name,
+            extra_env=strands_env,
+        )
+        assert result1.returncode == 0, f"Run 1 failed:\n{result1.stderr}"
+        wait_for_session_complete_sync(mgr, session_name, timeout=30)
+
+        result2 = _run_agent(
+            "examples/zero-code-examples/strands/run.py",
+            otlp_port,
+            session_name,
+            extra_env=strands_env,
+        )
+        assert result2.returncode == 0, f"Run 2 failed:\n{result2.stderr}"
+        wait_for_session_complete_sync(mgr, f"{session_name}-2", timeout=30)
+
+        s1 = mgr.sessions[session_name]
+        s2 = mgr.sessions[f"{session_name}-2"]
+
+        assert s1.is_complete and s2.is_complete
+        assert len(s1.spans) > 0
+        assert len(s2.spans) > 0
+        assert s1.trace_ids.isdisjoint(s2.trace_ids), (
+            f"Sessions share trace_ids: {s1.trace_ids & s2.trace_ids}"
+        )
+
+    def test_langchain_rerun_creates_separate_sessions(self, live_servers):
+        """Run the LangChain agent twice with the same session_name.
+        Each run must produce its own session."""
+        main_port, otlp_port, mgr = live_servers
+        session_name = "e2e-langchain-rerun"
+
+        result1 = _run_agent(
+            "examples/zero-code-examples/langchain/run.py",
+            otlp_port,
+            session_name,
+        )
+        assert result1.returncode == 0, f"Run 1 failed:\n{result1.stderr}"
+        wait_for_session_complete_sync(mgr, session_name, timeout=30)
+
+        result2 = _run_agent(
+            "examples/zero-code-examples/langchain/run.py",
+            otlp_port,
+            session_name,
+        )
+        assert result2.returncode == 0, f"Run 2 failed:\n{result2.stderr}"
+        wait_for_session_complete_sync(mgr, f"{session_name}-2", timeout=30)
+
+        s1 = mgr.sessions[session_name]
+        s2 = mgr.sessions[f"{session_name}-2"]
+
+        assert s1.is_complete and s2.is_complete
+        assert len(s1.spans) > 0
+        assert len(s2.spans) > 0
+        assert s1.trace_ids.isdisjoint(s2.trace_ids), (
+            f"Sessions share trace_ids: {s1.trace_ids & s2.trace_ids}"
+        )
+
+    def test_rerun_sessions_visible_via_api(self, live_servers):
+        """Both rerun sessions are visible in the API response."""
+        main_port, otlp_port, mgr = live_servers
+        session_name = "e2e-strands-rerun-api"
+        strands_env = {
+            "OTEL_SEMCONV_STABILITY_OPT_IN": "gen_ai_latest_experimental",
+        }
+
+        for run_idx in range(2):
+            result = _run_agent(
+                "examples/zero-code-examples/strands/run.py",
+                otlp_port,
+                session_name,
+                extra_env=strands_env,
+            )
+            assert result.returncode == 0, f"Run {run_idx + 1} failed:\n{result.stderr}"
+            expected_id = session_name if run_idx == 0 else f"{session_name}-2"
+            wait_for_session_complete_sync(mgr, expected_id, timeout=30)
+
+        resp = httpx.get(f"http://127.0.0.1:{main_port}/api/streaming/sessions")
+        assert resp.status_code == 200
+        session_ids = [s["sessionId"] for s in resp.json()["data"]]
+        assert session_name in session_ids
+        assert f"{session_name}-2" in session_ids

tests/integration/test_session_grouping.py

@@ -166,7 +166,7 @@ async def test_sequential_runs_unique_ids(self, trace_manager, otlp_client):
         await send_traces(otlp_client, body)
         await wait_for_session_complete(trace_manager, "repeated")
 
-        # Second run — same name but first is complete
+        # Second run — same name but new trace_id → new session
         body = make_trace_request(
             trace_id="run2",
             session_name="repeated",
@@ -317,3 +317,170 @@ async def test_span_limit_enforcement(self, trace_manager, otlp_client):
         await send_traces(otlp_client, body2)
 
         assert len(session.spans) == MAX_SPANS_PER_SESSION
+
+
+class TestSplitBatchReopen:
+    """Regression tests for the split-batch scenario where the OTLP
+    BatchSpanProcessor flushes one trace's spans across the session
+    completion boundary. Some child spans arrive before the grace period
+    fires, the root span arrives after. Because the trace_id was already
+    registered in the session, the session is reopened.
+
+    Bug report: strands-zero-code vs strands-zero-code-2 — a 3-turn
+    conversation was broken into 2 sessions because turn 3's spans
+    were split across the completion boundary."""
+
+    async def test_split_trace_reopens_completed_session(
+        self, trace_manager, otlp_client
+    ):
+        """When child spans arrive before completion and the root span
+        arrives after, the session reopens because the trace_id already
+        exists in the session."""
+        session_name = "split-reopen"
+
+        # Batch 1: turn 1 root + turn 2 child spans in same flush
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="sr-t1", session_name=session_name,
+            spans=[
+                make_genai_span(trace_id="sr-t1", span_id="t1-root", parent_span_id=None),
+                make_genai_span(trace_id="sr-t2", span_id="t2-child",
+                                parent_span_id="t2-root"),
+            ],
+        ))
+        await wait_for_session_complete(trace_manager, session_name)
+
+        session = trace_manager.sessions[session_name]
+        assert session.is_complete
+        assert "sr-t2" in session.trace_ids
+
+        # Batch 2: turn 2 root span arrives after completion
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="sr-t2", session_name=session_name,
+            spans=[
+                make_genai_span(trace_id="sr-t2", span_id="t2-root", parent_span_id=None),
+            ],
+        ))
+
+        assert not session.is_complete
+        assert len(trace_manager.sessions) == 1
+        assert session.trace_ids == {"sr-t1", "sr-t2"}
+
+        await wait_for_session_complete(trace_manager, session_name)
+        assert session.is_complete
+        assert len(session.spans) == 3
+
+    async def test_strands_three_turn_bug_repro(
+        self, trace_manager, otlp_client
+    ):
+        """Reproduces the exact bug from the Strands SDK report: the
+        BatchSpanProcessor flushes turns 1-2 and partial turn 3 in one
+        batch, then the rest of turn 3 in a second batch after the
+        session completes. All spans must end up in a single session."""
+        session_name = "strands-repro"
+
+        # Batch 1: turns 1 & 2 fully, plus turn 3 child spans
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="t1", session_name=session_name,
+            spans=[
+                make_genai_span(trace_id="t1", span_id="t1-llm", parent_span_id="t1-root"),
+                make_genai_span(trace_id="t1", span_id="t1-root", parent_span_id=None,
+                                name="invoke_agent"),
+                make_genai_span(trace_id="t2", span_id="t2-llm", parent_span_id="t2-root"),
+                make_genai_span(trace_id="t2", span_id="t2-tool",
+                                parent_span_id="t2-root", name="execute_tool roll_die"),
+                make_genai_span(trace_id="t2", span_id="t2-root", parent_span_id=None,
+                                name="invoke_agent"),
+                # Turn 3 child spans — flushed in same batch
+                make_genai_span(trace_id="t3", span_id="t3-llm", parent_span_id="t3-root"),
+                make_genai_span(trace_id="t3", span_id="t3-tool",
+                                parent_span_id="t3-root", name="execute_tool check_prime"),
+            ],
+        ))
+        await wait_for_session_complete(trace_manager, session_name)
+
+        session = trace_manager.sessions[session_name]
+        assert session.is_complete
+        assert "t3" in session.trace_ids
+
+        # Batch 2: turn 3 root span + remaining spans (after completion)
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="t3", session_name=session_name,
+            spans=[
+                make_genai_span(trace_id="t3", span_id="t3-loop",
+                                parent_span_id="t3-root", name="execute_event_loop_cycle"),
+                make_genai_span(trace_id="t3", span_id="t3-root", parent_span_id=None,
+                                name="invoke_agent"),
+            ],
+        ))
+
+        assert not session.is_complete
+        await wait_for_session_complete(trace_manager, session_name)
+
+        assert len(trace_manager.sessions) == 1
+        assert session.trace_ids == {"t1", "t2", "t3"}
+        assert len(session.spans) == 9
+
+    async def test_new_trace_after_completion_creates_new_session(
+        self, trace_manager, otlp_client
+    ):
+        """A completely new trace_id after session completion creates a
+        new session (not a reopen). This is the re-run case."""
+        session_name = "no-reopen"
+
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="run-1", session_name=session_name,
+            spans=[make_genai_span(trace_id="run-1", parent_span_id=None)],
+        ))
+        await wait_for_session_complete(trace_manager, session_name)
+
+        # New trace_id (not seen before) → new session
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="run-2", session_name=session_name,
+            spans=[make_genai_span(trace_id="run-2", parent_span_id=None)],
+        ))
+        await wait_for_session_complete(trace_manager, f"{session_name}-2")
+
+        assert len(trace_manager.sessions) == 2
+        assert trace_manager.sessions[session_name].trace_ids == {"run-1"}
+        assert trace_manager.sessions[f"{session_name}-2"].trace_ids == {"run-2"}
+
+    async def test_reopen_preserves_existing_spans_and_logs(
+        self, trace_manager, otlp_client
+    ):
+        """Reopening a session preserves all previously collected spans and logs."""
+        session_name = "preserve"
+
+        # Send spans and logs with trace_id "pres-t1", PLUS a child span
+        # from "pres-t2" so its trace_id is registered for reopen
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="pres-t1", session_name=session_name,
+            spans=[
+                make_genai_span(trace_id="pres-t1", parent_span_id=None),
+                make_genai_span(trace_id="pres-t2", span_id="t2-child",
+                                parent_span_id="t2-root"),
+            ],
+        ))
+        await send_logs(otlp_client, make_log_request(
+            trace_id="pres-t1", session_name=session_name,
+            log_records=[
+                make_genai_log("gen_ai.user.message", "Turn 1", trace_id="pres-t1"),
+            ],
+        ))
+        await wait_for_session_complete(trace_manager, session_name)
+
+        session = trace_manager.sessions[session_name]
+        spans_before = len(session.spans)
+        logs_before = len(session.logs)
+        assert spans_before >= 2
+        assert logs_before >= 1
+
+        # Reopen via split-batch trace_id match
+        await send_traces(otlp_client, make_trace_request(
+            trace_id="pres-t2", session_name=session_name,
+            spans=[make_genai_span(trace_id="pres-t2", span_id="t2-root",
+                                   parent_span_id=None)],
+        ))
+        await wait_for_session_complete(trace_manager, session_name)
+
+        assert len(session.spans) == spans_before + 1
+        assert len(session.logs) == logs_before