fix: Report per-index consumed capacity for writes (excludes transact/batch APIs for now)

crates/core/src/types/capacity.rs

@@ -193,6 +193,65 @@ impl ConsumedCapacity {
             local_secondary_indexes: None,
         }
     }
+
+    /// Build a write `ConsumedCapacity` whose aggregate total includes the base
+    /// table plus every affected secondary index.
+    ///
+    /// `base_cu` is the base-table write capacity. `gsi`/`lsi` map each affected
+    /// index name to its write capacity. DynamoDB reports the aggregate total as
+    /// `base + Σ(GSI) + Σ(LSI)`, and — in `INDEXES` mode — the per-table `Table`
+    /// capacity plus the two per-index maps.
+    ///
+    /// When `breakdown` is true (`INDEXES` mode) the `Table`,
+    /// `GlobalSecondaryIndexes` and `LocalSecondaryIndexes` fields are populated;
+    /// when false (`TOTAL` mode) only the aggregate is returned, but that
+    /// aggregate still reflects the index writes.
+    #[must_use]
+    pub fn write_indexed(
+        table_name: &str,
+        base_cu: f64,
+        gsi: HashMap<String, f64>,
+        lsi: HashMap<String, f64>,
+        breakdown: bool,
+    ) -> Self {
+        let total = base_cu + gsi.values().sum::<f64>() + lsi.values().sum::<f64>();
+        Self {
+            table_name: table_name.to_owned(),
+            capacity_units: total,
+            read_capacity_units: None,
+            write_capacity_units: Some(total),
+            table: breakdown.then(|| Capacity::write_units(base_cu)),
+            global_secondary_indexes: if breakdown { map_or_none(gsi) } else { None },
+            local_secondary_indexes: if breakdown { map_or_none(lsi) } else { None },
+        }
+    }
+}
+
+impl Capacity {
+    /// A write-only `Capacity` with the given units.
+    #[must_use]
+    pub fn write_units(cu: f64) -> Self {
+        Self {
+            capacity_units: cu,
+            read_capacity_units: None,
+            write_capacity_units: Some(cu),
+        }
+    }
+}
+
+/// Convert a per-index units map into `Some(map of Capacity)`, or `None` when
+/// empty so the field is omitted from the response.
+fn map_or_none(units: HashMap<String, f64>) -> Option<HashMap<String, Capacity>> {
+    if units.is_empty() {
+        None
+    } else {
+        Some(
+            units
+                .into_iter()
+                .map(|(name, cu)| (name, Capacity::write_units(cu)))
+                .collect(),
+        )
+    }
 }
 
 impl ItemCollectionMetrics {

crates/engine/src/batch_write_item.rs

@@ -199,6 +199,9 @@ pub async fn handle_batch_write_item(
         }
     }
 
+    // NOTE: per-index (INDEXES) breakdown for batch writes is deferred to the
+    // storage-layer capacity-reporting follow-up (Delete requests lack the old
+    // item needed to attribute per-index capacity). Base-table aggregate only.
     let consumed_capacity = capacity_helpers::batch_write_capacity(
         input.return_consumed_capacity,
         per_table_wcu.iter().map(|(t, cu)| (t.as_str(), *cu)),

crates/engine/src/capacity_helpers.rs

@@ -10,8 +10,8 @@
 use std::sync::atomic::AtomicU64;
 
 use extenddb_core::types::{
-    ConsumedCapacity, Item, ItemCollectionMetrics, KeySchemaElement, ReturnConsumedCapacity,
-    ReturnItemCollectionMetrics,
+    ConsumedCapacity, Item, ItemCollectionMetrics, KeySchemaElement, Projection, ProjectionType,
+    ReturnConsumedCapacity, ReturnItemCollectionMetrics, TableDescription, item_size_bytes,
 };
 
 /// Global counter for requests that used approximate consumed capacity.
@@ -52,6 +52,133 @@ pub fn write_capacity(
     }
 }
 
+/// Build a write `ConsumedCapacity` with a per-index breakdown for `INDEXES`
+/// mode, or the plain table-level capacity otherwise.
+///
+/// `base_cu` is the base-table write capacity; `item` is the item being written
+/// (used to determine sparse index membership and per-index projected size);
+/// `desc` is the table description (source of GSI/LSI definitions). When the
+/// caller does not request `INDEXES`, `desc` is unused and this behaves exactly
+/// like [`write_capacity`].
+#[must_use]
+pub fn write_capacity_indexed(
+    rcc: ReturnConsumedCapacity,
+    table_name: &str,
+    base_cu: f64,
+    item: &Item,
+    desc: &TableDescription,
+) -> Option<ConsumedCapacity> {
+    match rcc {
+        ReturnConsumedCapacity::None => None,
+        rcc => {
+            CAPACITY_REQUEST_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+            let (gsi, lsi) = index_write_units(item, desc);
+            let breakdown = rcc == ReturnConsumedCapacity::Indexes;
+            Some(ConsumedCapacity::write_indexed(
+                table_name, base_cu, gsi, lsi, breakdown,
+            ))
+        }
+    }
+}
+
+/// Compute per-GSI and per-LSI write capacity units contributed by `item`.
+///
+/// Returns `(gsi_units, lsi_units)` keyed by index name. An index is charged
+/// only when the item projects into it — i.e. the item contains every one of
+/// the index's key attributes (sparse-index semantics). Per-index capacity is
+/// `ceil(projected_item_size / 1KB)`, where the projected item contains only
+/// the attributes the index materializes (per its projection type).
+#[must_use]
+pub fn index_write_units(
+    item: &Item,
+    desc: &TableDescription,
+) -> (
+    std::collections::HashMap<String, f64>,
+    std::collections::HashMap<String, f64>,
+) {
+    let base_keys: Vec<&str> = desc
+        .key_schema
+        .iter()
+        .map(|k| k.attribute_name.as_str())
+        .collect();
+
+    let mut gsi = std::collections::HashMap::new();
+    if let Some(gsis) = &desc.global_secondary_indexes {
+        for g in gsis {
+            if let Some(cu) = one_index_write_units(item, &g.key_schema, &base_keys, &g.projection)
+            {
+                gsi.insert(g.index_name.clone(), cu);
+            }
+        }
+    }
+
+    let mut lsi = std::collections::HashMap::new();
+    if let Some(lsis) = &desc.local_secondary_indexes {
+        for l in lsis {
+            if let Some(cu) = one_index_write_units(item, &l.key_schema, &base_keys, &l.projection)
+            {
+                lsi.insert(l.index_name.clone(), cu);
+            }
+        }
+    }
+
+    (gsi, lsi)
+}
+
+/// Write units for a single index, or `None` if the item is not projected into
+/// it (missing an index key attribute — sparse index).
+fn one_index_write_units(
+    item: &Item,
+    index_key_schema: &[KeySchemaElement],
+    base_keys: &[&str],
+    projection: &Projection,
+) -> Option<f64> {
+    for ks in index_key_schema {
+        if !item.contains_key(&ks.attribute_name) {
+            return None;
+        }
+    }
+    let projected = project_index_item(item, index_key_schema, base_keys, projection);
+    Some(write_capacity_units(item_size_bytes(&projected)))
+}
+
+/// Build the subset of `item` that an index materializes, per its projection.
+fn project_index_item(
+    item: &Item,
+    index_key_schema: &[KeySchemaElement],
+    base_keys: &[&str],
+    projection: &Projection,
+) -> Item {
+    match projection.projection_type {
+        // ALL projects the entire item.
+        ProjectionType::All => item.clone(),
+        // KEYS_ONLY and INCLUDE always project index keys + base table keys.
+        ProjectionType::KeysOnly | ProjectionType::Include => {
+            let mut out = Item::new();
+            for ks in index_key_schema {
+                if let Some(v) = item.get(&ks.attribute_name) {
+                    out.insert(ks.attribute_name.clone(), v.clone());
+                }
+            }
+            for k in base_keys {
+                if let Some(v) = item.get(*k) {
+                    out.insert((*k).to_owned(), v.clone());
+                }
+            }
+            if projection.projection_type == ProjectionType::Include
+                && let Some(non_key) = &projection.non_key_attributes
+            {
+                for a in non_key {
+                    if let Some(v) = item.get(a) {
+                        out.insert(a.clone(), v.clone());
+                    }
+                }
+            }
+            out
+        }
+    }
+}
+
 /// Build a `Vec<ConsumedCapacity>` for a batch/transaction read, or `None` if not requested.
 /// One entry per distinct table name with real CU values.
 #[must_use]
@@ -74,7 +201,7 @@ pub fn batch_read_capacity<'a>(
 }
 
 /// Build a `Vec<ConsumedCapacity>` for a batch/transaction write, or `None` if not requested.
-/// One entry per distinct table name with real CU values.
+/// One entry per distinct table name with real CU values (base-table aggregate only).
 #[must_use]
 pub fn batch_write_capacity<'a>(
     rcc: ReturnConsumedCapacity,

crates/engine/src/delete_item.rs

@@ -104,13 +104,18 @@ pub async fn handle_delete_item(
         region: ctx.region.clone(),
     });
     let need_old_for_stream = stream.is_some();
+    // Consumed capacity — when requested, the total includes base + affected
+    // GSIs; INDEXES additionally returns the per-index breakdown. The deleted
+    // (old) item is needed to determine which indexes the delete propagated to.
+    let need_old_for_capacity =
+        input.return_consumed_capacity != extenddb_core::types::ReturnConsumedCapacity::None;
 
     let old_item = ctx
         .storage
         .delete_item(
             &key_info,
             &input.key,
-            return_old || need_old_for_stream,
+            return_old || need_old_for_stream || need_old_for_capacity,
             condition.as_ref(),
             &maps,
             stream.as_ref(),
@@ -129,13 +134,37 @@ pub async fn handle_delete_item(
             .map_or_else(|| item_size_bytes(&input.key), item_size_bytes),
     );
 
-    let output = DeleteItemOutput {
-        attributes: if return_old { old_item } else { None },
-        consumed_capacity: capacity_helpers::write_capacity(
+    // Consumed capacity — INDEXES mode uses the deleted (old) item to determine
+    // which indexes the delete propagated to, via one describe_table read. If
+    // nothing was deleted, only base-table capacity is reported.
+    let consumed_capacity = if input.return_consumed_capacity
+        != extenddb_core::types::ReturnConsumedCapacity::None
+        && let Some(cc_item) = old_item.as_ref()
+    {
+        let desc = ctx
+            .storage
+            .describe_table(
+                &ctx.account_id,
+                extenddb_core::types::DescribeTableInput {
+                    table_name: input.table_name.clone(),
+                },
+            )
+            .await
+            .map_err(storage_err_to_dynamo)?;
+        capacity_helpers::write_capacity_indexed(
             input.return_consumed_capacity,
             &input.table_name,
             wcu,
-        ),
+            cc_item,
+            &desc,
+        )
+    } else {
+        capacity_helpers::write_capacity(input.return_consumed_capacity, &input.table_name, wcu)
+    };
+
+    let output = DeleteItemOutput {
+        attributes: if return_old { old_item } else { None },
+        consumed_capacity,
         item_collection_metrics: capacity_helpers::item_metrics(
             input.return_item_collection_metrics,
             &key_info.key_schema,

crates/engine/src/put_item.rs

@@ -147,6 +147,33 @@ pub async fn handle_put_item(
     let item_bytes = item_size_bytes(&input.item);
     let wcu = capacity_helpers::write_capacity_units(item_bytes);
 
+    // Consumed capacity is computed before `input.item` is moved into storage.
+    // INDEXES mode additionally needs the item and index metadata to build the
+    // per-index breakdown, which requires one describe_table catalog read (only
+    // taken when INDEXES is requested).
+    let consumed_capacity =
+        if input.return_consumed_capacity != extenddb_core::types::ReturnConsumedCapacity::None {
+            let desc = ctx
+                .storage
+                .describe_table(
+                    &ctx.account_id,
+                    extenddb_core::types::DescribeTableInput {
+                        table_name: input.table_name.clone(),
+                    },
+                )
+                .await
+                .map_err(storage_err_to_dynamo)?;
+            capacity_helpers::write_capacity_indexed(
+                input.return_consumed_capacity,
+                &input.table_name,
+                wcu,
+                &input.item,
+                &desc,
+            )
+        } else {
+            capacity_helpers::write_capacity(input.return_consumed_capacity, &input.table_name, wcu)
+        };
+
     // Extract item collection metrics before item is moved into storage.
     let icm = capacity_helpers::item_metrics(
         input.return_item_collection_metrics,
@@ -186,11 +213,7 @@ pub async fn handle_put_item(
 
     let output = PutItemOutput {
         attributes: if return_old { old_item } else { None },
-        consumed_capacity: capacity_helpers::write_capacity(
-            input.return_consumed_capacity,
-            &input.table_name,
-            wcu,
-        ),
+        consumed_capacity,
         item_collection_metrics: icm,
     };
     let body = serialize_output(&output)?;

crates/engine/src/transact_write_items.rs

@@ -164,6 +164,11 @@ pub async fn handle_transact_write_items(
         })
         .sum();
 
+    // NOTE: per-index (INDEXES) consumed-capacity breakdown for transactions is
+    // deferred to the storage-layer capacity-reporting follow-up: the engine has
+    // no resulting/old item for Update/Delete sub-ops, but DynamoDB reports a
+    // per-index breakdown for those ops on GSI tables. Base-table aggregate
+    // only for now.
     let consumed_capacity = capacity_helpers::batch_write_capacity(
         input.return_consumed_capacity,
         per_table_wcu.iter().map(|(t, cu)| (t.as_str(), *cu)),

crates/engine/src/update_item.rs

@@ -242,6 +242,34 @@ pub async fn handle_update_item(
     let new_bytes = new_item.as_ref().map_or(0, item_size_bytes);
     let wcu = capacity_helpers::write_capacity_units(old_bytes.max(new_bytes));
 
+    // Consumed capacity — computed before old/new items are moved below.
+    // INDEXES mode uses the resulting item (new, or old for pure deletes) to
+    // determine sparse index membership, via one describe_table read.
+    let consumed_capacity = if input.return_consumed_capacity
+        != extenddb_core::types::ReturnConsumedCapacity::None
+        && let Some(cc_item) = new_item.as_ref().or(old_item.as_ref())
+    {
+        let desc = ctx
+            .storage
+            .describe_table(
+                &ctx.account_id,
+                extenddb_core::types::DescribeTableInput {
+                    table_name: input.table_name.clone(),
+                },
+            )
+            .await
+            .map_err(storage_err_to_dynamo)?;
+        capacity_helpers::write_capacity_indexed(
+            input.return_consumed_capacity,
+            &input.table_name,
+            wcu,
+            cc_item,
+            &desc,
+        )
+    } else {
+        capacity_helpers::write_capacity(input.return_consumed_capacity, &input.table_name, wcu)
+    };
+
     // Select the appropriate return value.
     // UPDATED_OLD and UPDATED_NEW return only the attributes that were
     // targeted by the update expression (top-level attribute of each action path).
@@ -259,11 +287,7 @@ pub async fn handle_update_item(
 
     let output = UpdateItemOutput {
         attributes,
-        consumed_capacity: capacity_helpers::write_capacity(
-            input.return_consumed_capacity,
-            &input.table_name,
-            wcu,
-        ),
+        consumed_capacity,
         item_collection_metrics: capacity_helpers::item_metrics(
             input.return_item_collection_metrics,
             &key_info.key_schema,