void onPartitionsAssigned(Collection<TopicPartition> partitions);
void onPartitionsRevoked(Collection<TopicPartition> partitions);

// Today
class MyListener implements ConsumerRebalanceListener {
    private final Consumer<?, ?> consumer;
    MyListener(Consumer<?, ?> consumer) { this.consumer = consumer; }
    public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
        consumer.commitSync(); // unrestricted reference
    }
    public void onPartitionsAssigned(Collection<TopicPartition> partitions) {}
}

// With this KIP
consumer.setRebalanceListener(new RebalanceListener() {
    // rest of the methods
    @Override
    public void onPartitionsRevoked(
            Collection<TopicPartition> partitions, RebalanceConsumer consumer) {
        consumer.commitSync();
    }
	@Override
    public void onPartitionsAssigned(
            Collection<TopicPartition> partitions, RebalanceConsumer consumer) {}
	
});
consumer.subscribe(topics);

package org.apache.kafka.clients.consumer;

/**
 * reuse Javadoc from ConsumerRebalanceListener for class javadoc and methodlevel javadoc.
 */
public interface RebalanceListener {
    /** 
     * +inherit Javadoc from ConsumerRebalanceListener. 
     * 
     * Consumer-aware variant of {@link ConsumerRebalanceListener#onPartitionsAssigned(Collection)}.
     *
     * <p>This method receives a {@link RebalanceConsumer}, a restricted view of the {@link Consumer}
     * that exposes only operations safe during a rebalance (e.g. seeking to externally stored offsets,
     * pausing partitions until caches warm up). The {@code RebalanceConsumer} is valid only for the
     * duration of this callback; storing a reference and using it later will throw
     * {@link IllegalStateException}.
     */
    void onPartitionsAssigned(Collection<TopicPartition> partitions, RebalanceConsumer rebalanceConsumer);

    //** 
     * +inherit Javadoc from ConsumerRebalanceListener. 
     * 
     * Consumer-aware variant of {@link ConsumerRebalanceListener#onPartitionsRevoked(Collection)}.
     *
     * <p>This method receives a {@link RebalanceConsumer}, a restricted view of the {@link Consumer}
     * that exposes only operations safe during a rebalance (e.g. seeking to externally stored offsets,
     * pausing partitions until caches warm up). The {@code RebalanceConsumer} is valid only for the
     * duration of this callback; storing a reference and using it later will throw
     * {@link IllegalStateException}.
     */
     void onPartitionsRevoked(Collection<TopicPartition> partitions, RebalanceConsumer rebalanceConsumer);

    /**
     * Consumer-aware variant of {@link #onPartitionsLost(Collection)}.
     *
     * <p>This method receives a {@link RebalanceConsumer}, a restricted view of the {@link Consumer}
     * that exposes only operations safe during a rebalance. The {@code RebalanceConsumer} is valid
     * only for the duration of this callback; storing a reference and using it later will throw
     * {@link IllegalStateException}.
     *
     * <h4>Important difference from {@code onPartitionsRevoked}</h4>
     *
     * <p>Unlike {@link #onPartitionsRevoked(Collection, RebalanceConsumer)}, where the affected
     * partitions are still part of the consumer's assignment until the callback completes, the
     * partitions passed to this method have <em>already been removed</em> from the assignment
     * before this callback fires. This means:
     *
     * <ul>
     *   <li>{@link RebalanceConsumer#commitSync() commitSync()}/{@link RebalanceConsumer#commitAsync() commitAsync()}
     *       will not commit the current assignment, since the consumer no longer includes the lost partitions.</li>
     *   <li>{@link RebalanceConsumer#commitSync(java.util.Map) commitSync(offsets)}
     *       /{@link RebalanceConsumer#commitAsync(java.util.Map, OffsetCommitCallback) commitAsync(offsets)}
     *       with the lost partitions explicitly will be rejected by the broker since the consumer
     *       no longer owns them.</li>
     *   <li>{@link RebalanceConsumer#position(org.apache.kafka.common.TopicPartition) position()} and
     *       {@link RebalanceConsumer#seek(org.apache.kafka.common.TopicPartition, long) seek()} will
     *       fail for the lost partitions since they are not in the assignment.</li>
     * </ul>
     *
     * <p>The {@link RebalanceConsumer} does not add guardrails against these operations on lost
     * partitions. Users should be aware that offset management for lost partitions is not
     * possible in this callback.
     *
     * @param partitions The list of partitions that were assigned to the consumer and now have been reassigned
     *                   to other consumers. With both the Classic and Consumer protocols, this will always include
     *                   all partitions that were previously assigned to the consumer.
     * @param consumer   A restricted view of the {@link Consumer} that is only valid for the duration of this callback.
     *                   See {@link RebalanceConsumer} for the set of permitted operations.
     * @throws org.apache.kafka.common.errors.WakeupException    If raised from a nested call to {@link RebalanceConsumer}
     * @throws org.apache.kafka.common.errors.InterruptException If raised from a nested call to {@link RebalanceConsumer}
     *
     * @see #onPartitionsLost(Collection)
     * @see RebalanceConsumer
     */
    @Override
    default void onPartitionsLost(Collection<TopicPartition> partitions, RebalanceConsumer rebalanceConsumer) {
        onPartitionsRevoked(partitions, rebalanceConsumer);
    }
}

package org.apache.kafka.clients.consumer;

/**
 * {@inheritDoc}
 * @deprecated
 * @see RebalanceListener
 */
@Deprecated(since = "4.4", forRemoval = true)
public interface ConsumerRebalanceListener extends RebalanceListener {
   
    // Other methods stay unchanged.

	/**
     * Called after partitions have been assigned to this consumer.
     *
     * <p>The default implementation delegates to
     * {@link #onPartitionsAssigned(Collection)}. Override this method
     * instead of the one-arg variant when the callback needs to interact
     * with the consumer (commit offsets, seek, pause/resume).
     *
     * <p>When this method is overridden, the one-arg
     * {@code onPartitionsAssigned(Collection)} is not called unless the
     * override explicitly invokes it.
     *
     * @param partitions the newly assigned partitions
     * @param consumer   a restricted view of the consumer; supports seek,
     *                   commit, pause/resume, and read-only queries
     */
    @Override
    default void onPartitionsAssigned(Collection<TopicPartition> partitions, RebalanceConsumer consumer) {
        onPartitionsAssigned(partitions);
    }

    /**
     * Called before partitions are revoked from this consumer during a
     * cooperative rebalance, or before all partitions are revoked during
     * an eager rebalance.
     *
     * <p>The default implementation delegates to
     * {@link #onPartitionsRevoked(Collection)}.
     *
     * @param partitions the partitions being revoked
     * @param consumer   a restricted view of the consumer
     */
    @Override
    default void onPartitionsRevoked(Collection<TopicPartition> partitions, RebalanceConsumer consumer) {
        onPartitionsRevoked(partitions);
    }

    /**
     * Called when partitions have been lost. Unlike revocation, the consumer
     * no longer owns these partitions and commits will fail.
     *
     * <p>The default implementation delegates to {@link #onPartitionsLost(Collection)}.
     *
     * @param partitions the partitions that were lost
     * @param consumer   a restricted view of the consumer
     */
    @Override
    default void onPartitionsLost(Collection<TopicPartition> partitions, RebalanceConsumer consumer) {
        onPartitionsLost(partitions);
    }
}

package org.apache.kafka.clients.consumer;

/**
 * A restricted view of a {@link Consumer} passed to {@link ConsumerRebalanceListener} methods.
 * This object is only valid for the duration of the rebalance callback. Any use outside of this scope,
 * will result in an {@link IllegalStateException}.
 *
 * <p>This interface exposes only operations that are safe and meaningful
 * during partition rebalancing:
 *
 * <ul>
 *   <li><b>Offset management:</b> {@code commitSync}, {@code commitAsync},
 *       {@code committed}, {@code position}</li>
 *   <li><b>Seek:</b> {@code seek}, {@code seekToBeginning},
 *       {@code seekToEnd}</li>
 *   <li><b>Partition state:</b> {@code assignment}, {@code paused},
 *       {@code pause}, {@code resume}</li>
 *   <li><b>Read-only queries:</b> {@code partitionsFor}, {@code listTopics},
 *       {@code offsetsForTimes}, {@code beginningOffsets},
 *       {@code endOffsets}, {@code metrics}, {@code clientInstanceId}, {@code currentLag}, {@code groupMetadata}</li>
 * </ul>
 *
 * <p>Operations that are dangerous, inapplicable during a rebalance are not present
 * on this interface:
 *
 * <ul>
 *   <li>{@code poll()} - would cause re-entrant polling</li>
 *   <li>{@code close()} - would terminate the consumer mid-rebalance</li>
 *   <li>{@code subscribe()}, {@code unsubscribe()} - would corrupt
 *       subscription state during rebalance</li>
 *   <li>{@code registerMetricForSubscription}, {@code unregisterMetricForSubscription} - metric
 *       registration does not belong during rebalance events</li>
 *   <li>{@code assign()} - conflicts with group-managed assignment</li>
 *   <li>{@code wakeup()} - would interrupt the rebalance</li>
 *   <li>{@code enforceRebalance()} - would trigger re-entrant rebalance</li>
 * </ul>
 */
public interface RebalanceConsumer {

    // --- Offset management ---

    void commitSync();
    void commitSync(Duration timeout);
    void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets);
    void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets, Duration timeout);
    void commitAsync();
    void commitAsync(OffsetCommitCallback callback);
    void commitAsync(Map<TopicPartition, OffsetAndMetadata> offsets,
                     OffsetCommitCallback callback);

    Map<TopicPartition, OffsetAndMetadata> committed(Set<TopicPartition> partitions);
    Map<TopicPartition, OffsetAndMetadata> committed(Set<TopicPartition> partitions,
                                                      Duration timeout);
    long position(TopicPartition partition);
    long position(TopicPartition partition, Duration timeout);

    // --- Seek ---

    void seek(TopicPartition partition, long offset);
    void seek(TopicPartition partition, OffsetAndMetadata offsetAndMetadata);
    void seekToBeginning(Collection<TopicPartition> partitions);
    void seekToEnd(Collection<TopicPartition> partitions);

    // --- Partition state ---

    Set<TopicPartition> assignment();
    void pause(Collection<TopicPartition> partitions);
    void resume(Collection<TopicPartition> partitions);
    Set<TopicPartition> paused();

    // --- Read-only queries ---

	Uuid clientInstanceId(Duration timeout);
    Map<TopicPartition, Long> beginningOffsets(Collection<TopicPartition> partitions);
    Map<TopicPartition, Long> beginningOffsets(Collection<TopicPartition> partitions,
                                                Duration timeout);
    Map<TopicPartition, Long> endOffsets(Collection<TopicPartition> partitions);
    Map<TopicPartition, Long> endOffsets(Collection<TopicPartition> partitions,
                                         Duration timeout);
    Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(
            Map<TopicPartition, Long> timestampsToSearch);
    Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(
            Map<TopicPartition, Long> timestampsToSearch, Duration timeout);
    List<PartitionInfo> partitionsFor(String topic);
    List<PartitionInfo> partitionsFor(String topic, Duration timeout);
    Map<String, List<PartitionInfo>> listTopics();
    Map<String, List<PartitionInfo>> listTopics(Duration timeout);
	OptionalLong currentLag(TopicPartition topicPartition);
    ConsumerGroupMetadata groupMetadata();
}

/** Update class javadoc with setRebalanceListener pattern. */

/** Also add a reference to {@link #setRebalanceListener} on existing subscribe() methods, to help with 5.0 migration. */

/**
 * Register a {@link ConsumerRebalanceListener} to be invoked when the consumer's partition assignment
 * changes during a rebalance. This is the preferred way to register a rebalance listener, replacing
 * the previous pattern of passing a listener to the {@code subscribe()} methods (which are now deprecated).
 *
 * <p>
 * The listener can be updated at any time. The change takes effect on the next {@link #poll(Duration)}
 * invocation: any rebalance triggered during that poll (or subsequent polls) will use the most recently
 * set listener.
 *
 * <p>
 * Passing {@code null} removes the current listener, so that no callbacks are invoked on future rebalances.
 *
 * <p>
 * If a listener was previously registered via one of the deprecated {@code subscribe(topics, listener)}
 * variants, calling this method will override it. 
 * 
 * <p>
 * The same is true for multiple registrations via this method. The final value that was set before the next
 * {@code poll()} will receive callback invocations during rebalance events.
 *
 * @param callback The listener to invoke on partition assignment changes, or {@code null} to remove
 *                 the current listener.
 *
 * @see ConsumerRebalanceListener
 * @see RebalanceConsumer
 */
void setRebalanceListener(RebalanceListener callback);

/** Use {@link #setRebalanceListener} instead. */
@Deprecated(since = "4.4", forRemoval = true)
void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener callback);

/** Use {@link #setRebalanceListener} instead. */
@Deprecated(since = "4.4", forRemoval = true)
void subscribe(Collection<String> topics, ConsumerRebalanceListener callback);

/** Use {@link #setRebalanceListener} instead. */
@Deprecated(since = "4.4", forRemoval = true)
void subscribe(Pattern pattern, ConsumerRebalanceListener callback);

// Before 
listener.onPartitionsAssigned(partitions); 
listener.onPartitionsRevoked(partitions); 
listener.onPartitionsLost(partitions);

// After
try (var ignored = new DelegatingRebalanceConsumer(this.consumer)) {
    listener.onPartitionsAssigned(partitions, view);
}

// This will throw IllegalStateException when someApplicationMethod() is called
private RebalanceConsumer storedView;

void onPartitionsAssigned(Collection<TopicPartition> partitions,
                          RebalanceConsumer consumer) {
    this.storedView = consumer; // Stores reference
}

void someApplicationMethod() {
    storedView.commitSync(); // throws IllegalStateException: view expired
}