feat: add _partitions_observed tracking via observe() for in_progress/pending distinction

- Add _partitions_observed: set[str] to track partitions where observe() has been called
- Update state property to emit 4 fields: in_progress, completed, expected, is_partition_discovery_complete
- num_partitions_in_progress = len(observed) - completed (worker started but not finished)
- Dropped num_partitions_not_started (derivable from expected - in_progress - completed)
- _cleanup_if_done adds partition to observed set to prevent negative in_progress
- Updated all test assertions for new semantics

Author: devin-ai-integration[bot]

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -190,6 +190,9 @@ def __init__(
         # (_set_initial_state does not read them back). On resume, they reset to 0.
         self._num_partitions_completed: int = 0
         self._is_partition_discovery_complete: bool = False
+        # Tracks partition keys for which observe() has been called (worker produced at least one record).
+        # Only len() is used in state emission; the set itself is never serialized.
+        self._partitions_observed: set[str] = set()
 
         self._set_initial_state(stream_state)
 
@@ -223,9 +226,9 @@ def state(self) -> MutableMapping[str, Any]:
             state["lookback_window"] = self._lookback_window
         if self._parent_state is not None:
             state["parent_state"] = self._parent_state
+        num_observed = len(self._partitions_observed)
         state["partitioned_stream_status"] = {
-            "num_partitions_in_progress": self._generated_partitions_count
-            - self._num_partitions_completed,
+            "num_partitions_in_progress": num_observed - self._num_partitions_completed,
             "num_partitions_completed": self._num_partitions_completed,
             "num_partitions_expected": self._generated_partitions_count,
             "is_partition_discovery_complete": self._is_partition_discovery_complete,
@@ -552,11 +555,11 @@ def observe(self, record: Record) -> None:
             return
 
         self._synced_some_data = True
+        partition_key = self._to_partition_key(record.associated_slice.partition)
+        self._partitions_observed.add(partition_key)
         self._update_global_cursor(record_cursor)
         if not self._use_global_cursor:
-            self._cursor_per_partition[
-                self._to_partition_key(record.associated_slice.partition)
-            ].observe(record)
+            self._cursor_per_partition[partition_key].observe(record)
 
     def _update_global_cursor(self, value: Any) -> None:
         if (
@@ -581,6 +584,8 @@ def _cleanup_if_done(self, partition_key: str) -> None:
 
         seq = self._partition_key_to_index.pop(partition_key)
         self._processing_partitions_indexes.remove(seq)
+        # Ensure completed partitions are counted as observed (handles partitions with no records)
+        self._partitions_observed.add(partition_key)
         self._num_partitions_completed += 1
 
         logger.debug(f"Partition {partition_key} fully processed and cleaned up.")

unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py

@@ -405,9 +405,11 @@ def run_mocked_test(
         assert "num_partitions_expected" in partitioned_status
         assert "is_partition_discovery_complete" in partitioned_status
         assert partitioned_status["num_partitions_in_progress"] >= 0
+        assert partitioned_status["num_partitions_completed"] >= 0
         assert (
             partitioned_status["num_partitions_expected"]
-            >= partitioned_status["num_partitions_completed"]
+            >= partitioned_status["num_partitions_in_progress"]
+            + partitioned_status["num_partitions_completed"]
         )
         _strip_partitioned_stream_status(final_state_dict)
         assert final_state_dict == expected_state
@@ -3695,6 +3697,12 @@ def test_given_no_partitions_processed_when_close_partition_then_no_state_update
     assert partitioned_status["num_partitions_completed"] == 0
     assert partitioned_status["num_partitions_expected"] == 0
     assert partitioned_status["is_partition_discovery_complete"] is True
+    # Invariant: in_progress + completed <= expected
+    assert (
+        partitioned_status["num_partitions_in_progress"]
+        + partitioned_status["num_partitions_completed"]
+        <= partitioned_status["num_partitions_expected"]
+    )
     assert state == {
         "use_global_cursor": False,
         "lookback_window": 0,
@@ -3785,7 +3793,8 @@ def test_given_unfinished_first_parent_partition_no_parent_state_update():
     state = cursor.state
     partitioned_status = state.pop("partitioned_stream_status", None)
     assert partitioned_status is not None
-    assert partitioned_status["num_partitions_in_progress"] == 1
+    # observe() not called in this test, so in_progress comes only from _cleanup_if_done adding to observed
+    assert partitioned_status["num_partitions_in_progress"] == 0
     assert partitioned_status["num_partitions_completed"] == 1
     assert partitioned_status["num_partitions_expected"] == 2
     assert partitioned_status["is_partition_discovery_complete"] is True
@@ -3887,7 +3896,8 @@ def test_given_unfinished_last_parent_partition_with_partial_parent_state_update
     state = cursor.state
     partitioned_status = state.pop("partitioned_stream_status", None)
     assert partitioned_status is not None
-    assert partitioned_status["num_partitions_in_progress"] == 1
+    # observe() not called in this test, so in_progress comes only from _cleanup_if_done adding to observed
+    assert partitioned_status["num_partitions_in_progress"] == 0
     assert partitioned_status["num_partitions_completed"] == 1
     assert partitioned_status["num_partitions_expected"] == 2
     assert partitioned_status["is_partition_discovery_complete"] is True