fix(plugins): add opt-in exactly-once delivery to BigQueryAgentAnalyticsPlugin

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -213,6 +213,8 @@ async def wrapper(
 
 # gRPC Error Codes
 _GRPC_DEADLINE_EXCEEDED = 4
+_GRPC_ALREADY_EXISTS = 6
+_GRPC_OUT_OF_RANGE = 11
 _GRPC_INTERNAL = 13
 _GRPC_UNAVAILABLE = 14
 
@@ -1754,6 +1756,9 @@ class BigQueryLoggerConfig:
         emit the final answer via a dedicated tool (e.g.
         ``submit_final_response``) rather than a plain-text final event. Empty
         (the default) preserves today's behavior.
+      exactly_once_delivery: When ``True``, writes use a per-loop
+        ``COMMITTED`` stream with offset tracking so an ambiguous retry
+        cannot write duplicate rows. ``False`` (default) keeps ``_default``.
   """
 
   enabled: bool = True
@@ -1829,6 +1834,9 @@ class BigQueryLoggerConfig:
   # behavior.
   final_response_tool_names: frozenset[str] = frozenset()
 
+  # Off by default; see class docstring for tradeoffs.
+  exactly_once_delivery: bool = False
+
 
 # ==============================================================================
 # HELPER: TRACE MANAGER (Async-Safe with ContextVars)
@@ -2188,6 +2196,8 @@ def __init__(
       retry_config: RetryConfig,
       queue_max_size: int,
       shutdown_timeout: float,
+      exactly_once_delivery: bool = False,
+      create_stream: Optional[Callable[[], Coroutine[Any, Any, str]]] = None,
   ):
     """Initializes the instance.
 
@@ -2200,6 +2210,10 @@ def __init__(
         retry_config: Retry configuration.
         queue_max_size: Max size of the in-memory queue.
         shutdown_timeout: Max time to wait for shutdown.
+        exactly_once_delivery: If ``True``, ``write_stream`` must be a
+          dedicated ``COMMITTED`` stream; appends carry a tracked offset.
+        create_stream: Factory for a fresh ``COMMITTED`` stream, used to
+          rotate away from a desynced stream. ``None`` disables rotation.
     """
     self.write_client = write_client
     self.arrow_schema = arrow_schema
@@ -2208,6 +2222,13 @@ def __init__(
     self.flush_interval = flush_interval
     self.retry_config = retry_config
     self.shutdown_timeout = shutdown_timeout
+    self.exactly_once_delivery = exactly_once_delivery
+    self.create_stream = create_stream
+    # All four only used when exactly_once_delivery is True.
+    self._next_offset = 0
+    self._offset_desynced = False
+    self._stream_finalized = False
+    self._rotation_backoff_until = 0.0
 
     self._visual_builder = _is_visual_builder.get()
 
@@ -2235,6 +2256,7 @@ def __init__(
         "unexpected_error": 0,
         "shutdown_timeout": 0,
         "shutdown_cancelled": 0,
+        "offset_conflict": 0,
     }
 
   async def flush(self) -> None:
@@ -2279,6 +2301,8 @@ def get_drop_stats(self) -> dict[str, int]:
       ``shutdown_timeout``: rows still queued when shutdown timed out.
       ``shutdown_cancelled``: rows still queued when shutdown was
         cancelled from outside (e.g. a host close timeout).
+      ``offset_conflict``: (exactly-once only) offset tracking is desynced
+        from the server; batches drop until stream rotation succeeds.
 
     Returns:
         A copy of the per-reason drop counters.
@@ -2432,12 +2456,69 @@ async def _batch_writer(self) -> None:
         else:
           break
 
+  def _confirm_delivered(self, offset: Optional[int], row_count: int) -> None:
+    # Assignment, not `+=`: a late timeout plus an ALREADY_EXISTS retry can
+    # confirm the same batch twice; both must resolve to the same next offset.
+    if offset is not None:
+      self._next_offset = offset + row_count
+
+  def _mark_offset_ambiguous(self) -> None:
+    # The batch may have landed despite the failure; reusing its offset could
+    # get ALREADY_EXISTS and silently swallow a future batch. Poison instead.
+    if self.exactly_once_delivery and not self._offset_desynced:
+      self._offset_desynced = True
+      logger.error(
+          "Ambiguous write failure on stream %s (Data Loss): offset tracking"
+          " is no longer trustworthy; batches will be dropped until rotation"
+          " to a fresh stream succeeds.",
+          self.write_stream,
+      )
+
+  async def _rotate_stream(self) -> bool:
+    # Swaps a desynced stream for a fresh one so delivery can resume; backs
+    # off between attempts to avoid hammering CreateWriteStream.
+    if (
+        self.create_stream is None
+        or time.monotonic() < self._rotation_backoff_until
+    ):
+      return False
+    await self._finalize_stream()
+    try:
+      new_stream = await self.create_stream()
+    except Exception as e:
+      self._rotation_backoff_until = time.monotonic() + 30.0
+      logger.warning("Failed to rotate desynced write stream: %s", e)
+      return False
+    self.write_stream = new_stream
+    self._next_offset = 0
+    self._offset_desynced = False
+    self._stream_finalized = False
+    logger.warning(
+        "Rotated to fresh write stream %s after offset desync; resuming"
+        " exactly-once writes.",
+        new_stream,
+    )
+    return True
+
   async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None:
     """Writes a batch of rows to BigQuery with retry logic.
 
     Args:
         rows: list of row dictionaries to write.
     """
+    if self.exactly_once_delivery and self._offset_desynced:
+      if not await self._rotate_stream():
+        # Desync was already logged at state entry; drop until a later
+        # rotation attempt succeeds.
+        self._dropped["offset_conflict"] += len(rows)
+        logger.debug(
+            "Dropping batch: offset tracking for stream %s is desynced."
+            " Total rows dropped (offset conflict): %s",
+            self.write_stream,
+            self._dropped["offset_conflict"],
+        )
+        return
+
     attempt = 0
     delay = self.retry_config.initial_delay
 
@@ -2458,6 +2539,13 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None:
       )
       req.arrow_rows.writer_schema.serialized_schema = serialized_schema
       req.arrow_rows.rows.serialized_record_batch = serialized_batch
+      # Frozen for this batch's whole retry sequence; resolved only through
+      # _confirm_delivered.
+      offset_for_batch = (
+          self._next_offset if self.exactly_once_delivery else None
+      )
+      if offset_for_batch is not None:
+        req.offset = offset_for_batch
     except Exception as e:
       self._dropped["arrow_prep_failed"] += len(rows)
       logger.error(
@@ -2470,12 +2558,17 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None:
       return
 
     while attempt <= self.retry_config.max_retries:
+      request_sent = False
+      in_band_rejection = False
       try:
 
         async def requests_iter() -> AsyncIterator[Any]:
+          nonlocal request_sent
+          request_sent = True
           yield req
 
         async def perform_write() -> None:
+          nonlocal in_band_rejection
           # The AppendRows streaming RPC does not auto-populate the
           # request-routing header, so writes to any region other than
           # the US multiregion fail with a "session not found" /
@@ -2496,16 +2589,56 @@ async def perform_write() -> None:
             error_code = getattr(error, "code", None)
             if error_code and error_code != 0:
               error_message = getattr(error, "message", "Unknown error")
+
+              if (
+                  offset_for_batch is not None
+                  and error_code == _GRPC_ALREADY_EXISTS
+              ):
+                # Already applied by a prior attempt whose ack we missed;
+                # confirmed delivered, not an error.
+                logger.info(
+                    "BigQuery append at offset %s for stream %s was already"
+                    " applied; retry confirmed as a no-op (no duplicate"
+                    " written).",
+                    offset_for_batch,
+                    self.write_stream,
+                )
+                self._confirm_delivered(offset_for_batch, len(rows))
+                return
+
+              if (
+                  offset_for_batch is not None
+                  and error_code == _GRPC_OUT_OF_RANGE
+              ):
+                # Local offset desynced from the server's; drop rather than
+                # guess a corrective offset, and stop trusting this stream.
+                self._offset_desynced = True
+                self._dropped["offset_conflict"] += len(rows)
+                logger.error(
+                    "BigQuery Storage Write offset conflict at offset %s for"
+                    " stream %s (Data Loss): local offset tracking is out of"
+                    " sync with the server. Total rows dropped (offset"
+                    " conflict): %s",
+                    offset_for_batch,
+                    self.write_stream,
+                    self._dropped["offset_conflict"],
+                )
+                return
+
               logger.warning(
                   "BigQuery Write API returned error code %s: %s",
                   error_code,
                   error_message,
               )
+
               if error_code in [
                   _GRPC_DEADLINE_EXCEEDED,
                   _GRPC_INTERNAL,
                   _GRPC_UNAVAILABLE,
               ]:
+                # An in-band error means the rows were not appended, so the
+                # frozen offset stays free to reuse on retry.
+                in_band_rejection = True
                 raise ServiceUnavailable(error_message)
 
               if "schema mismatch" in error_message.lower():
@@ -2526,7 +2659,14 @@ async def perform_write() -> None:
                 )
               self._dropped["non_retryable"] += len(rows)
               return
-          return
+            # One response per request; return without waiting for the
+            # stream close, which could time out and force a spurious retry.
+            self._confirm_delivered(offset_for_batch, len(rows))
+            return
+          # Zero responses: this batch's fate is unknown; refuse to confirm.
+          if offset_for_batch is not None:
+            self._dropped["unexpected_error"] += len(rows)
+            self._mark_offset_ambiguous()
 
         await asyncio.wait_for(perform_write(), timeout=30.0)
         return
@@ -2547,6 +2687,10 @@ async def perform_write() -> None:
               e,
               self._dropped["retry_exhausted"],
           )
+          # A never-sent request or an in-band rejection is definitively
+          # not applied, so only an ack-lost failure is ambiguous.
+          if request_sent and not in_band_rejection:
+            self._mark_offset_ambiguous()
           return
 
         sleep_time = min(
@@ -2570,8 +2714,26 @@ async def perform_write() -> None:
             self._dropped["unexpected_error"],
             exc_info=True,
         )
+        if request_sent:
+          self._mark_offset_ambiguous()
         return
 
+  async def _finalize_stream(self, timeout: float = 5.0) -> None:
+    # Best-effort, attempted once; rows in a COMMITTED stream are already
+    # visible, so failures here lose nothing.
+    if not self.exactly_once_delivery or self._stream_finalized:
+      return
+    self._stream_finalized = True
+    try:
+      await asyncio.wait_for(
+          self.write_client.finalize_write_stream(name=self.write_stream),
+          timeout=min(timeout, 5.0),
+      )
+    except Exception as e:
+      logger.warning(
+          "Failed to finalize write stream %s: %s", self.write_stream, e
+      )
+
   def _drain_queue_and_count(self, reason: str) -> int:
     """Counts and discards everything still queued (loss accounting)."""
     drained = 0
@@ -2597,6 +2759,7 @@ async def shutdown(self, timeout: float = 5.0) -> None:
         timeout: Maximum time to wait for the queue to drain.
     """
     self._shutdown = True
+    finalize_deadline = time.monotonic() + timeout
     logger.info("BatchProcessor shutting down, draining queue...")
 
     # Signal the writer to wake up and check shutdown status
@@ -2664,13 +2827,16 @@ async def shutdown(self, timeout: float = 5.0) -> None:
         raise
       except Exception as e:
         logger.error("Error during BatchProcessor shutdown: %s", e)
+    # Spend only what is left of the caller's timeout budget.
+    await self._finalize_stream(max(0.0, finalize_deadline - time.monotonic()))
 
   async def close(self) -> None:
     """Closes the processor and flushes remaining items."""
     if self._shutdown:
       return
 
     self._shutdown = True
+    finalize_deadline = time.monotonic() + self.shutdown_timeout
     # Wait for queue to be empty
     try:
       await asyncio.wait_for(self._queue.join(), timeout=self.shutdown_timeout)
@@ -2700,6 +2866,7 @@ async def close(self) -> None:
     if drained:
       self._dropped["shutdown_timeout"] += drained
       logger.warning("%d queued row(s) dropped by close timeout.", drained)
+    await self._finalize_stream(max(0.0, finalize_deadline - time.monotonic()))
 
 
 # ==============================================================================
@@ -4092,6 +4259,23 @@ def _format_content_safely(
       logger.warning("Content formatter failed: %s", e)
       return "[FORMATTING FAILED]", False
 
+  def _table_resource_path(self) -> str:
+    """The table's BigQuery Storage API resource path (no stream suffix)."""
+    return f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}"
+
+  async def _create_committed_write_stream(
+      self, write_client: "BigQueryWriteAsyncClient"
+  ) -> str:
+    """Creates the dedicated COMMITTED stream that, unlike ``_default``,
+    accepts append offsets, and returns its full resource name."""
+    stream = await write_client.create_write_stream(
+        parent=self._table_resource_path(),
+        write_stream=bq_storage_types.WriteStream(
+            type_=bq_storage_types.WriteStream.Type.COMMITTED
+        ),
+    )
+    return stream.name
+
   async def _close_write_transport(self, write_client: Any) -> None:
     """Best-effort bounded close for a BigQuery write-client transport."""
     transport = getattr(write_client, "transport", None)
@@ -4243,19 +4427,40 @@ def get_credentials() -> google.auth.credentials.Credentials:
           client_options=options,
       )
 
-      if not self._write_stream_name:
-        self._write_stream_name = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default"
+      if self.config.exactly_once_delivery:
+        # COMMITTED streams need a single writer, so each loop gets its own.
+        try:
+          write_stream_name = await self._create_committed_write_stream(
+              write_client
+          )
+        except BaseException:
+          # Same ownership rule as below: the client exists but no _LoopState
+          # can own it yet, so close it here rather than leak the transport.
+          await self._close_write_transport(write_client)
+          raise
+      else:
+        if not self._write_stream_name:
+          self._write_stream_name = f"{self._table_resource_path()}/_default"
+        write_stream_name = self._write_stream_name
 
       try:
         batch_processor = BatchProcessor(
             write_client=write_client,
             arrow_schema=self.arrow_schema,
-            write_stream=self._write_stream_name,
+            write_stream=write_stream_name,
             batch_size=self.config.batch_size,
             flush_interval=self.config.batch_flush_interval,
             retry_config=self.config.retry_config,
             queue_max_size=self.config.queue_max_size,
             shutdown_timeout=self.config.shutdown_timeout,
+            exactly_once_delivery=self.config.exactly_once_delivery,
+            create_stream=(
+                functools.partial(
+                    self._create_committed_write_stream, write_client
+                )
+                if self.config.exactly_once_delivery
+                else None
+            ),
         )
       except BaseException:
         # The write client already exists but no _LoopState can own it yet.