fix(telemetry): preserve span kinds in SQLite traces

src/google/adk/telemetry/sqlite_span_exporter.py

@@ -30,6 +30,7 @@
 from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.sdk.trace.export import SpanExportResult
 from opentelemetry.trace import SpanContext
+from opentelemetry.trace import SpanKind
 from opentelemetry.trace import TraceFlags
 from opentelemetry.trace import TraceState
 from opentelemetry.util.types import AttributeValue
@@ -46,7 +47,8 @@
   end_time_unix_nano INTEGER,
   session_id TEXT,
   invocation_id TEXT,
-  attributes_json TEXT
+  attributes_json TEXT,
+  kind TEXT NOT NULL DEFAULT 'INTERNAL'
 );
 """
 
@@ -68,8 +70,9 @@
   end_time_unix_nano,
   session_id,
   invocation_id,
-  attributes_json
-) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?);
+  attributes_json,
+  kind
+) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
 """
 
 _DEFAULT_TIMEOUT_SECONDS = 30.0
@@ -80,6 +83,9 @@ class SqliteSpanExporter(SpanExporter):
 
   This is intended for local development (e.g. `adk web`) to allow reloading
   traces for older sessions after process restart.
+
+  Span kinds survive export and reload. Existing databases are upgraded in
+  place; spans saved before kind persistence retain the INTERNAL default.
   """
 
   def __init__(self, *, db_path: str):
@@ -101,10 +107,19 @@ def _get_connection(self) -> sqlite3.Connection:
   def _ensure_schema(self) -> None:
     with self._lock:
       conn = self._get_connection()
-      conn.execute(_CREATE_SPANS_TABLE)
-      conn.execute(_CREATE_SESSION_INDEX)
-      conn.execute(_CREATE_TRACE_INDEX)
-      conn.commit()
+      with conn:
+        conn.execute("BEGIN IMMEDIATE")
+        conn.execute(_CREATE_SPANS_TABLE)
+        if not any(
+            row["name"] == "kind"
+            for row in conn.execute("PRAGMA table_info(spans)")
+        ):
+          conn.execute(
+              "ALTER TABLE spans ADD COLUMN kind TEXT NOT NULL"
+              " DEFAULT 'INTERNAL'"
+          )
+        conn.execute(_CREATE_SESSION_INDEX)
+        conn.execute(_CREATE_TRACE_INDEX)
 
   def _serialize_attributes(
       self, attributes: Mapping[str, AttributeValue]
@@ -167,6 +182,7 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
               session_id,
               invocation_id,
               self._serialize_attributes(attributes),
+              span.kind.name,
           ))
         conn.executemany(_INSERT_SPAN, rows)
         conn.commit()
@@ -221,6 +237,7 @@ def _row_to_readable_span(self, row: sqlite3.Row) -> ReadableSpan:
         name=row["name"] or "",
         context=context,
         parent=parent,
+        kind=SpanKind[row["kind"]],
         attributes=attributes,
         start_time=row["start_time_unix_nano"],
         end_time=row["end_time_unix_nano"],

tests/unittests/telemetry/test_sqlite_span_exporter.py

@@ -16,11 +16,13 @@
 
 import json
 from pathlib import Path
+import sqlite3
 
 from google.adk.telemetry.sqlite_span_exporter import SqliteSpanExporter
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
 from opentelemetry.trace import SpanContext
+from opentelemetry.trace import SpanKind
 from opentelemetry.trace import TraceFlags
 from opentelemetry.trace import TraceState
 
@@ -32,6 +34,7 @@ def _create_span(
     parent_span_id: int | None = None,
     name: str = "test_span",
     attributes: dict | None = None,
+    kind: SpanKind = SpanKind.INTERNAL,
     start_time: int = 1000,
     end_time: int = 2000,
 ) -> ReadableSpan:
@@ -58,6 +61,7 @@ def _create_span(
       name=name,
       context=context,
       parent=parent,
+      kind=kind,
       attributes=attributes or {},
       start_time=start_time,
       end_time=end_time,
@@ -227,6 +231,86 @@ def test_round_trip_preserves_span_attributes(tmp_path):
   assert retrieved.attributes == original_attributes
 
 
+def test_span_kinds_survive_exporter_restart(tmp_path):
+  db_path = tmp_path / "test.db"
+  kinds = [
+      SpanKind.CLIENT,
+      SpanKind.SERVER,
+      SpanKind.PRODUCER,
+      SpanKind.CONSUMER,
+  ]
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+  spans = [
+      _create_span(
+          span_id=index + 1,
+          kind=kind,
+          attributes={"gen_ai.conversation.id": "session-kind"},
+          start_time=index,
+      )
+      for index, kind in enumerate(kinds)
+  ]
+  assert exporter.export(spans) == SpanExportResult.SUCCESS
+  exporter.shutdown()
+
+  reopened = SqliteSpanExporter(db_path=str(db_path))
+  try:
+    restored = reopened.get_all_spans_for_session("session-kind")
+    assert [(span.context.span_id, span.kind) for span in restored] == [
+        (index + 1, kind) for index, kind in enumerate(kinds)
+    ]
+  finally:
+    reopened.shutdown()
+
+
+def test_legacy_database_retains_spans_and_accepts_new_kinds(tmp_path):
+  db_path = tmp_path / "legacy.db"
+  with sqlite3.connect(db_path) as conn:
+    conn.execute("""
+        CREATE TABLE spans (
+          span_id TEXT PRIMARY KEY, trace_id TEXT NOT NULL,
+          parent_span_id TEXT, name TEXT NOT NULL,
+          start_time_unix_nano INTEGER, end_time_unix_nano INTEGER,
+          session_id TEXT, invocation_id TEXT, attributes_json TEXT
+        )
+    """)
+    conn.execute(
+        "INSERT INTO spans VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+        (
+            "0000000000000001",
+            "000000000000000000000000000def45",
+            None,
+            "legacy_operation",
+            1000,
+            2000,
+            "session-legacy",
+            None,
+            '{"gen_ai.conversation.id": "session-legacy"}',
+        ),
+    )
+
+  exporter = SqliteSpanExporter(db_path=str(db_path))
+  try:
+    assert (
+        exporter.export([
+            _create_span(
+                span_id=2,
+                name="new_operation",
+                kind=SpanKind.CLIENT,
+                attributes={"gen_ai.conversation.id": "session-legacy"},
+                start_time=3000,
+            )
+        ])
+        == SpanExportResult.SUCCESS
+    )
+    restored = exporter.get_all_spans_for_session("session-legacy")
+    assert [(span.name, span.kind) for span in restored] == [
+        ("legacy_operation", SpanKind.INTERNAL),
+        ("new_operation", SpanKind.CLIENT),
+    ]
+  finally:
+    exporter.shutdown()
+
+
 def test_spans_with_parent_context_exported_correctly(tmp_path):
   db_path = tmp_path / "test.db"
   exporter = SqliteSpanExporter(db_path=str(db_path))