google-cloud-storage 1.21.0

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::connector::{Connection, Connector};
use super::replay_buffer::{ReplayBuffer, ReplayChunk};
use super::{Client, MAX_WRITE_CHUNK_SIZE, TonicStreaming, is_finalized, persisted_size};
use crate::Error;
use crate::google::storage::v2::{
    BidiWriteObjectRequest, BidiWriteObjectResponse, bidi_write_object_request::Data,
};
use std::collections::VecDeque;
use std::sync::Arc;

use gaxi::grpc::from_status::to_gax_error;
use gaxi::grpc::tonic::Result as TonicResult;
use tokio::sync::mpsc::{Receiver, Sender};
use tokio::sync::oneshot;

type LoopResult<T> = std::result::Result<T, Error>;

/// The intent sent from the foreground task to the background worker.
pub enum UploadIntent {
    Append(BidiWriteObjectRequest),
    Flush(
        BidiWriteObjectRequest,
        oneshot::Sender<crate::Result<BidiWriteObjectResponse>>,
    ),
    Finalize(
        BidiWriteObjectRequest,
        oneshot::Sender<crate::Result<BidiWriteObjectResponse>>,
    ),
}

/// What [`Worker::run`] must do after handling one server message.
enum ResponseAction<S> {
    /// Keep using the current stream.
    Continue,
    /// The stream was replaced; adopt these halves.
    Reconnected(Connection<S>),
    /// The server closed the stream and the upload is complete.
    Finished,
}

/// Where a [`BidiWriteObjectResponse`] came from.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ResponseOrigin {
    /// Read from the live stream; one that matches nothing was unprompted.
    Stream,
    /// The `state_lookup` handshake of a reconnected stream; always expected.
    Handshake,
}

/// Which kind of confirmation a [`PendingRequest`] is waiting for.
#[derive(Debug)]
enum PendingKind {
    Flush,
    Finalize,
}

/// Tracks an in-flight flush or finalize request awaiting server confirmation.
#[derive(Debug)]
struct PendingRequest {
    kind: PendingKind,
    request: BidiWriteObjectRequest,
    sender: oneshot::Sender<crate::Result<BidiWriteObjectResponse>>,
}

impl PendingRequest {
    fn new(
        kind: PendingKind,
        request: BidiWriteObjectRequest,
        sender: oneshot::Sender<crate::Result<BidiWriteObjectResponse>>,
    ) -> Self {
        Self {
            kind,
            request,
            sender,
        }
    }

    /// Returns `true` when `response` confirms this request completed.
    ///
    /// Once `persisted_size` reaches `self.request.write_offset`, a [`PendingKind::Flush`] is
    /// satisfied by an unfinalized response and a [`PendingKind::Finalize`] is satisfied by a
    /// finalized response.
    fn is_satisfied(&self, response: &BidiWriteObjectResponse, persisted_size: i64) -> bool {
        if persisted_size < self.request.write_offset {
            return false;
        }
        match self.kind {
            PendingKind::Flush => !is_finalized(response),
            PendingKind::Finalize => is_finalized(response),
        }
    }

    fn complete(self, response: crate::Result<BidiWriteObjectResponse>) {
        let _ = self.sender.send(response);
    }
}

/// The background worker that manages the live gRPC stream, unacknowledged chunk replay, and
/// automatic reconnection.
pub struct Worker<C> {
    connector: Connector<C>,
    replay_buffer: ReplayBuffer,
    pending_requests: VecDeque<PendingRequest>,
    /// Tracks the highest byte offset acknowledged by the server.
    persisted_size: i64,
    /// Tracks if the client intends to complete the upload by sending a Finalize intent.
    finalized: bool,
    /// Tracks if a worker-initiated `state_lookup` request is awaiting a server response.
    ///
    /// The worker injects such a request when the replay buffer crosses its high watermark, so that
    /// a response is already owed by the time the buffer fills and `ack()` can drain it. Up to two
    /// more maximum-size appends may still arrive before that response does.
    self_flush_outstanding: bool,
}

impl<C> Worker<C> {
    pub fn new(connector: Connector<C>) -> Self {
        Self::with_replay_buffer(connector, ReplayBuffer::new())
    }

    /// Creates a [`Worker`] with a caller-provided [`ReplayBuffer`].
    ///
    /// The transport uses this to seed the buffer with the payload sent in the opening request, and
    /// tests use it to exercise the capacity limits without buffering
    /// [`DEFAULT_REPLAY_BUFFER_SIZE`][super::replay_buffer::DEFAULT_REPLAY_BUFFER_SIZE] bytes.
    pub fn with_replay_buffer(connector: Connector<C>, replay_buffer: ReplayBuffer) -> Self {
        Self {
            connector,
            replay_buffer,
            pending_requests: VecDeque::new(),
            persisted_size: 0,
            finalized: false,
            self_flush_outstanding: false,
        }
    }

    /// Sets the initial acknowledged byte offset from the opening or reopen response.
    pub fn with_persisted_size(mut self, persisted_size: i64) -> Self {
        self.persisted_size = persisted_size;
        self
    }

    /// Returns `true` if the replay buffer has crossed its high watermark and neither a user
    /// flush/finalize nor a worker-initiated `state_lookup` is already awaiting a response.
    ///
    /// Because plain appends do not elicit a server response, this keeps a response owed, either to
    /// a pending user flush/finalize or to a worker probe, whenever the buffer is above the
    /// watermark, so the loop stays live once [`ReplayBuffer::is_full`] pauses the intent channel.
    fn needs_watermark_flush(&self) -> bool {
        !self.self_flush_outstanding
            && self.pending_requests.is_empty()
            && self.replay_buffer.unpersisted_bytes()
                >= self
                    .replay_buffer
                    .capacity()
                    .saturating_sub(2 * MAX_WRITE_CHUNK_SIZE)
    }
}

impl<C> Worker<C>
where
    C: Client + Clone + 'static,
    <C as Client>::Stream: TonicStreaming,
{
    pub async fn run(
        mut self,
        connection: Connection<C::Stream>,
        mut requests: Receiver<UploadIntent>,
    ) -> LoopResult<()> {
        let (mut rx, mut tx) = (connection.rx, connection.tx);

        let error = loop {
            tokio::select! {
                m = rx.next_message() => {
                    match self.handle_response(m).await {
                        Err(e) => break Some(e),
                        // Successful end of stream, return without error.
                        Ok(ResponseAction::Finished) => break None,
                        // The stream was replaced, adopt the new halves.
                        Ok(ResponseAction::Reconnected(connection)) => {
                            (rx, tx) = (connection.rx, connection.tx);
                        }
                        // Message handled on the existing stream. Re-evaluate the watermark, since
                        // the ack may have left the buffer above it.
                        Ok(ResponseAction::Continue) => {
                            self.maybe_send_watermark_flush(&tx).await;
                        }
                    }
                },
                intent = requests.recv(), if !self.replay_buffer.is_full() => {
                    match intent {
                        Some(intent) => {
                            let request = self.process_intent(intent);
                            // If sending fails, `rx.next_message()` will observe the stream's
                            // status or closure on the next iteration and reconnect, replaying the
                            // chunk or pending request recorded by `process_intent`.
                            if let Err(e) = tx.send(request).await {
                                tracing::debug!(
                                    "error sending request on bidi write stream: {e:?}"
                                );
                            }
                        }
                        None => {
                            drop(tx);
                            break self.wait_for_server_completion(rx).await;
                        }
                    }
                }
            }
        };

        if let Some(e) = error {
            let shared_error = Arc::new(e);
            self.drain_intents_on_error(requests, Arc::clone(&shared_error))
                .await;
            return Err(Error::ser(shared_error));
        }

        Ok(())
    }

    fn process_intent(&mut self, intent: UploadIntent) -> BidiWriteObjectRequest {
        match intent {
            UploadIntent::Append(mut req) => {
                if let Some(Data::ChecksummedData(ref cd)) = req.data {
                    let crc32c = cd.crc32c.unwrap_or_else(|| crc32c::crc32c(&cd.content));
                    self.replay_buffer.push(ReplayChunk::new(
                        req.write_offset,
                        cd.content.clone(),
                        crc32c,
                    ));
                }
                // Piggyback a `state_lookup` on this append once the replay buffer crosses its high
                // watermark. Only `state_lookup` elicits a response, and only a response drives
                // `ReplayBuffer::ack()`.
                if self.needs_watermark_flush() {
                    req.flush = true;
                    req.state_lookup = true;
                    self.self_flush_outstanding = true;
                }
                req
            }
            UploadIntent::Flush(req, sender) => {
                assert!(
                    req.state_lookup,
                    "state_lookup must be true for Flush intents"
                );
                assert!(req.flush, "flush must be true for Flush intents");
                self.pending_requests.push_back(PendingRequest::new(
                    PendingKind::Flush,
                    req.clone(),
                    sender,
                ));
                req
            }
            UploadIntent::Finalize(req, sender) => {
                assert!(req.flush, "flush must be true for Finalize intents");
                assert!(
                    req.finish_write,
                    "finish_write must be true for Finalize intents"
                );
                self.finalized = true;
                self.pending_requests.push_back(PendingRequest::new(
                    PendingKind::Finalize,
                    req.clone(),
                    sender,
                ));
                req
            }
        }
    }

    /// Handles an incoming response message or stream completion from the server.
    ///
    /// Returns the action [`Self::run`] must take next, or an error if the failure is unrecoverable
    /// or reconnection failed.
    async fn handle_response(
        &mut self,
        message: TonicResult<Option<BidiWriteObjectResponse>>,
    ) -> LoopResult<ResponseAction<C::Stream>> {
        let response = match message {
            Ok(Some(msg)) => msg,
            Ok(None) => {
                // The server closed the stream. That is the expected end of a finalized upload
                // once the finalize response has been delivered. Otherwise a flush or finalize is
                // still waiting for its response, or the client never asked to finalize, so
                // reconnect rather than drop unacknowledged work and fail later client writes.
                if !self.pending_requests.is_empty() || !self.finalized {
                    let connection = self
                        .reconnect(Error::io("stream closed unexpectedly"))
                        .await?;
                    return Ok(ResponseAction::Reconnected(connection));
                }
                return Ok(ResponseAction::Finished);
            }
            Err(e) => {
                // The finalized object has already been handed to the caller; a trailing stream
                // error cannot change it. Mirror the `Ok(None)` arm above instead of reconnecting
                // against a finalized object.
                if self.finalized && self.pending_requests.is_empty() {
                    tracing::debug!("ignoring stream error after finalize completed: {e:?}");
                    return Ok(ResponseAction::Finished);
                }
                let connection = self.reconnect(to_gax_error(e)).await?;
                return Ok(ResponseAction::Reconnected(connection));
            }
        };
        self.handle_response_success(response, ResponseOrigin::Stream)?;
        Ok(ResponseAction::Continue)
    }

    /// Processes a successful [`BidiWriteObjectResponse`] from the server.
    ///
    /// Updates acknowledged offsets in the replay buffer and completes any matching in-flight flush
    /// or finalize requests. Returns an error if the server reports a `persisted_size` lower than
    /// previously acknowledged bytes.
    fn handle_response_success(
        &mut self,
        response: BidiWriteObjectResponse,
        origin: ResponseOrigin,
    ) -> LoopResult<()> {
        // Every response ends the wait for a worker-initiated `state_lookup`, including one that
        // carries no `write_status`. Clearing the flag only on the acknowledged-bytes path would
        // suppress all later watermark probes and stall the worker once the replay buffer fills.
        let self_flush_outstanding = std::mem::take(&mut self.self_flush_outstanding);

        // The service periodically refreshes the write handle, often in a response that carries
        // nothing else. Record it before the early return below, so a later `reconnect` opens its
        // stream with the latest handle. The handshake response was already recorded by
        // `Connector::connect_attempt`.
        if origin == ResponseOrigin::Stream {
            self.connector.handle_response(&response);
        }

        let Some(persisted_size) = persisted_size(&response) else {
            // A response with no `write_status` carries no progress; its write handle, if any,
            // was recorded above.
            tracing::debug!("Received BidiWriteObjectResponse with no write_status: {response:?}");
            return Ok(());
        };

        if persisted_size < self.persisted_size {
            return Err(Error::io(format!(
                "server persisted_size ({persisted_size}) regressed below \
                 acknowledged offset ({})",
                self.persisted_size
            )));
        }
        // TODO(#5716): Reject persisted_size exceeding the highest sent write offset to guard
        // against concurrent writers on the same object generation.

        self.persisted_size = persisted_size;
        self.replay_buffer.ack(persisted_size);

        // Pending requests are answered in order, so stop at the first one this response does not
        // satisfy. That one is still waiting for its own response, which the service owes us, so
        // the loop stays live.
        let mut matched = false;
        while let Some(front) = self.pending_requests.front()
            && front.is_satisfied(&response, persisted_size)
        {
            let Some(request) = self.pending_requests.pop_front() else {
                break;
            };
            request.complete(Ok(response.clone()));
            matched = true;
        }

        // An object reported as finalized before the client requested finalization (or before its
        // target offset was reached) cannot accept further writes or replayed chunks.
        if is_finalized(&response) && (!self.finalized || !self.pending_requests.is_empty()) {
            return Err(Error::io("object is already finalized"));
        }

        // A reconnect handshake and a worker-initiated `state_lookup` have no entry in
        // `pending_requests`, so their responses legitimately match nothing. Do not report them as
        // unprompted.
        if origin == ResponseOrigin::Stream && !matched && !self_flush_outstanding {
            tracing::debug!(
                "Received unprompted BidiWriteObjectResponse from server: {response:?}"
            );
        }
        Ok(())
    }

    /// Sends a standalone `flush + state_lookup` probe when the replay buffer remains above its
    /// watermark and no user `Flush` or `Finalize` is pending.
    async fn maybe_send_watermark_flush(&mut self, tx: &Sender<BidiWriteObjectRequest>) {
        if self.needs_watermark_flush()
            && let Some(write_offset) = self.replay_buffer.end_offset()
        {
            let request = BidiWriteObjectRequest {
                write_offset,
                flush: true,
                state_lookup: true,
                ..BidiWriteObjectRequest::default()
            };
            // If sending fails, the gRPC call has ended, and `rx.next_message()` will observe its
            // status or closure on a later iteration and reconnect. No probe is in flight in that
            // case, so leave `self_flush_outstanding` unset; `reconnect()` re-arms the probe after
            // replay if the buffer is still above the watermark.
            match tx.send(request).await {
                Ok(()) => self.self_flush_outstanding = true,
                Err(e) => {
                    tracing::debug!("error sending watermark probe on bidi write stream: {e:?}");
                }
            }
        }
    }

    /// Reconnects, then restores the session state on the new stream.
    ///
    /// Replays every unacknowledged chunk, re-sends the flush and finalize requests still waiting
    /// for a response, and re-arms the watermark probe if the replay buffer is still above its
    /// watermark. Returns an error if the connector gives up on reconnecting (the retry policy
    /// classifies `last_error` as permanent or the retry or redirect budget is exhausted), or if
    /// the handshake response reports a `persisted_size` regression or an unexpectedly finalized
    /// object.
    async fn reconnect(&mut self, last_error: Error) -> LoopResult<Connection<C::Stream>> {
        // Any response the worker-initiated `state_lookup` was waiting on will never arrive on the
        // dead stream.
        self.self_flush_outstanding = false;

        let (initial_response, connection) = self
            .connector
            .reconnect(last_error, self.persisted_size)
            .await?;

        // Process the reconnect handshake response. The spec is always `Append` here, so the new
        // stream opened with `state_lookup: true`. If the service finalized the object before the
        // previous stream broke, `initial_response` carries a finalized `Resource`
        // (`finalize_time.is_some()`) and satisfies the pending `Finalize` here, so it is not
        // re-sent. Otherwise the handshake reports progress as a bare `PersistedSize` or as an
        // unfinalized `Resource` (`finalize_time.is_none()`); neither satisfies a pending
        // `Finalize` (`PendingRequest::is_satisfied` requires a finalized response), so it is
        // re-sent below. A pending `Flush` whose `write_offset` is already persisted is completed
        // here instead of being re-sent.
        self.handle_response_success(initial_response, ResponseOrigin::Handshake)?;

        // Replay all unpersisted chunks, then re-send the pending flush / finalize requests. A
        // send fails only when the new stream's request channel is already closed, so stop at the
        // first failure: the next `rx.next_message()` poll in `run()` observes the stream's status
        // or closure and reconnects again, with the replay buffer and `pending_requests` intact.
        for chunk in self.replay_buffer.chunks_to_replay() {
            if let Err(e) = connection.tx.send(chunk.to_request()).await {
                tracing::debug!("error replaying chunk on bidi write stream: {e:?}");
                return Ok(connection);
            }
        }
        for pending in &self.pending_requests {
            if let Err(e) = connection.tx.send(pending.request.clone()).await {
                tracing::debug!("error re-sending pending request on bidi write stream: {e:?}");
                return Ok(connection);
            }
        }

        // Replayed chunks set neither `flush` nor `state_lookup`. If the buffer is still above the
        // watermark and no user request is pending, nothing would elicit a response on the new
        // stream, so restore the invariant explicitly.
        self.maybe_send_watermark_flush(&connection.tx).await;

        Ok(connection)
    }

    /// Drains the response stream after the intent channel closed and `tx` was dropped.
    ///
    /// Dropping `tx` half-closes the request stream, so the server sends any remaining responses
    /// and then ends the call. Unlike [`Self::run`], this does not reconnect on an error or early
    /// close. On the graceful path the transport closes the intent channel only after `close()`
    /// has flushed or `finalize()` has received its response, so every chunk is already
    /// acknowledged and no request is pending. If the writer is dropped without closing, any
    /// unacknowledged chunks are abandoned by design; nothing can ask for them again. Returns the
    /// error that ended the stream, if any.
    async fn wait_for_server_completion(&mut self, mut rx: C::Stream) -> Option<Error> {
        loop {
            match rx.next_message().await {
                Ok(Some(msg)) => {
                    if let Err(e) = self.handle_response_success(msg, ResponseOrigin::Stream) {
                        break Some(e);
                    }
                }
                Ok(None) => break None,
                // Same as in `handle_response`: a trailing error after the finalize response was
                // delivered cannot change the finalized object.
                Err(e) if self.finalized && self.pending_requests.is_empty() => {
                    tracing::debug!("ignoring stream error after finalize completed: {e:?}");
                    break None;
                }
                Err(e) => break Some(to_gax_error(e)),
            }
        }
    }

    async fn drain_intents_on_error(
        &mut self,
        mut requests: Receiver<UploadIntent>,
        shared_error: Arc<Error>,
    ) {
        for pending in self.pending_requests.drain(..) {
            pending.complete(Err(Error::ser(Arc::clone(&shared_error))));
        }
        // Drain remaining requests to notify pending flush/finalize intents if the stream failed.
        requests.close();
        while let Some(intent) = requests.recv().await {
            match intent {
                UploadIntent::Flush(_, sender) | UploadIntent::Finalize(_, sender) => {
                    let _ = sender.send(Err(Error::ser(Arc::clone(&shared_error))));
                }
                UploadIntent::Append(_) => {}
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::mocks::{MockTestClient, mock_connector};
    use super::super::replay_buffer::{DEFAULT_REPLAY_BUFFER_SIZE, MIN_REPLAY_BUFFER_SIZE};
    use super::super::state::AppendObjectSpecState;
    use super::super::tests::permanent_error;
    use super::*;
    use crate::google::storage::v2::{
        AppendObjectSpec, BidiWriteHandle, BidiWriteObjectRequest, BidiWriteObjectResponse,
        ChecksummedData, Object, bidi_write_object_request::FirstMessage,
        bidi_write_object_response::WriteStatus,
    };
    use gaxi::grpc::tonic::Response as TonicResponse;
    use gaxi::grpc::tonic::Result as TonicResult;
    use gaxi::grpc::tonic::Status;
    use tokio::sync::mpsc;
    use tokio::sync::oneshot;

    type TestWorkerContext = (
        tokio::task::JoinHandle<LoopResult<()>>,
        mpsc::Sender<UploadIntent>,
        mpsc::Receiver<BidiWriteObjectRequest>,
        mpsc::Sender<TonicResult<BidiWriteObjectResponse>>,
    );

    type ReconnectingWorkerContext = (
        tokio::task::JoinHandle<LoopResult<()>>,
        mpsc::Sender<UploadIntent>,
        mpsc::Receiver<BidiWriteObjectRequest>,
        mpsc::Sender<TonicResult<BidiWriteObjectResponse>>,
        mpsc::Receiver<mpsc::Receiver<BidiWriteObjectRequest>>,
        mpsc::Sender<TonicResult<BidiWriteObjectResponse>>,
    );

    /// Defines the payload size of the synthetic appends used in the watermark tests.
    const TEST_CHUNK_SIZE: usize = 512;

    /// Places the watermark (`capacity - 2 * MAX_WRITE_CHUNK_SIZE`) at exactly two test chunks, so
    /// the second append is the one that crosses it.
    const TEST_CAPACITY: usize = 2 * MAX_WRITE_CHUNK_SIZE + 2 * TEST_CHUNK_SIZE;

    /// The `AppendObjectSpec` that reconnecting test workers hand to the connector.
    fn test_spec_state() -> AppendObjectSpecState {
        AppendObjectSpecState::Append {
            spec: AppendObjectSpec {
                bucket: "projects/_/buckets/test-bucket".into(),
                object: "test-object".into(),
                ..Default::default()
            },
            initial_chunk: None,
        }
    }

    /// Wires `mock` into a connector and spawns a worker over fresh request, response, and intent
    /// channels. `spec_state`, when set, is what the connector uses on reconnect.
    fn spawn_worker(
        mock: MockTestClient,
        capacity: usize,
        spec_state: Option<AppendObjectSpecState>,
    ) -> TestWorkerContext {
        let (request_tx, request_rx) = mpsc::channel(10);
        let (response_tx, response_rx) = mpsc::channel(10);
        let (intent_tx, intent_rx) = mpsc::channel(10);
        let connection = Connection::new(request_tx, response_rx);

        let mut connector = mock_connector(mock);
        if let Some(state) = spec_state {
            connector.set_spec_state(state);
        }
        let worker = Worker::with_replay_buffer(connector, ReplayBuffer::with_capacity(capacity));
        let handle = tokio::spawn(worker.run(connection, intent_rx));

        (handle, intent_tx, request_rx, response_tx)
    }

    fn spawn_test_worker() -> TestWorkerContext {
        spawn_test_worker_with_capacity(DEFAULT_REPLAY_BUFFER_SIZE)
    }

    fn spawn_test_worker_with_capacity(capacity: usize) -> TestWorkerContext {
        let mut mock = MockTestClient::new();
        mock.expect_start().never();
        spawn_worker(mock, capacity, None)
    }

    fn append_intent(write_offset: i64, len: usize) -> UploadIntent {
        let content = bytes::Bytes::from(vec![b'x'; len]);
        let crc32c = crc32c::crc32c(&content);
        UploadIntent::Append(BidiWriteObjectRequest {
            write_offset,
            data: Some(Data::ChecksummedData(ChecksummedData {
                content,
                crc32c: Some(crc32c),
            })),
            ..Default::default()
        })
    }

    fn flush_intent(
        write_offset: i64,
    ) -> (
        UploadIntent,
        oneshot::Receiver<crate::Result<BidiWriteObjectResponse>>,
    ) {
        let (flush_tx, flush_rx) = oneshot::channel();
        let intent = UploadIntent::Flush(
            BidiWriteObjectRequest {
                write_offset,
                flush: true,
                state_lookup: true,
                ..Default::default()
            },
            flush_tx,
        );
        (intent, flush_rx)
    }

    fn persisted_size_response(size: i64) -> BidiWriteObjectResponse {
        BidiWriteObjectResponse {
            write_status: Some(WriteStatus::PersistedSize(size)),
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn run_append() -> anyhow::Result<()> {
        let (handle, tx, mut request_rx, _response_tx) = spawn_test_worker();

        let append_request = BidiWriteObjectRequest {
            write_offset: 10,
            ..Default::default()
        };
        tx.send(UploadIntent::Append(append_request)).await?;

        let stream_req = request_rx.recv().await.unwrap();
        assert_eq!(stream_req.write_offset, 10);

        drop(tx);
        tokio::task::yield_now().await;
        drop(_response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_flush() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) = spawn_test_worker();

        let (intent, mut flush_rx) = flush_intent(100);
        tx.send(intent).await?;

        let stream_req = request_rx.recv().await.unwrap();
        assert!(stream_req.flush);
        assert!(stream_req.state_lookup);

        // Act.
        // An ack below `target_offset` (50 < 100) leaves the flush pending.
        response_tx.send(Ok(persisted_size_response(50))).await?;
        tokio::task::yield_now().await;
        assert!(matches!(
            flush_rx.try_recv(),
            Err(oneshot::error::TryRecvError::Empty)
        ));

        let server_resp = persisted_size_response(100);
        response_tx.send(Ok(server_resp.clone())).await?;

        // Assert.
        let received_resp = flush_rx.await??;
        assert_eq!(received_resp.write_status, server_resp.write_status);

        drop(tx);
        tokio::task::yield_now().await;
        drop(response_tx);
        handle.await??;
        Ok(())
    }

    /// Builds the response the service sends once it finalized the object.
    fn finalized_response(size: i64) -> BidiWriteObjectResponse {
        BidiWriteObjectResponse {
            write_status: Some(WriteStatus::Resource(Object {
                name: "test-obj".into(),
                size,
                finalize_time: Some(prost_types::Timestamp::default()),
                ..Default::default()
            })),
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn run_finalize() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) = spawn_test_worker();

        let (intent, finalize_rx) = finalize_intent(100);
        tx.send(intent).await?;

        let stream_req = request_rx.recv().await.unwrap();
        assert!(stream_req.finish_write);

        // Act.
        let server_resp = finalized_response(100);
        response_tx.send(Ok(server_resp.clone())).await?;

        // Assert.
        let received_resp = finalize_rx.await??;
        assert_eq!(received_resp.write_status, server_resp.write_status);

        drop(response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_finalize_ignores_trailing_stream_error() -> anyhow::Result<()> {
        // Arrange.
        // `spawn_test_worker` sets `expect_start().never()`, so any reconnect attempt would panic.
        let (handle, tx, mut request_rx, response_tx) = spawn_test_worker();

        let (intent, finalize_rx) = finalize_intent(100);
        tx.send(intent).await?;
        let _ = request_rx.recv().await.unwrap();
        response_tx.send(Ok(finalized_response(100))).await?;
        let _ = finalize_rx.await??;

        // Act.
        // The intent channel is still open, so `handle_response` sees the error.
        response_tx
            .send(Err(Status::unavailable("connection reset")))
            .await?;

        // Assert.
        // The finalized object was already delivered, so the worker finishes cleanly.
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_finalize_ignores_trailing_stream_error_after_intents_close() -> anyhow::Result<()>
    {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) = spawn_test_worker();

        let (intent, finalize_rx) = finalize_intent(100);
        tx.send(intent).await?;
        let _ = request_rx.recv().await.unwrap();
        response_tx.send(Ok(finalized_response(100))).await?;
        let _ = finalize_rx.await??;

        // Close the intent channel first, so `wait_for_server_completion` sees the error.
        drop(tx);
        tokio::task::yield_now().await;

        // Act.
        response_tx
            .send(Err(Status::unavailable("connection reset")))
            .await?;

        // Assert.
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_finalize_ignores_resource_without_finalize_time() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) = spawn_test_worker();

        let (intent, mut finalize_rx) = finalize_intent(100);
        tx.send(intent).await?;
        let _ = request_rx.recv().await.unwrap();

        // Act.
        // A bare resource is what a create or handle-less takeover stream returns as its first
        // message. It reaches the target offset, but the object was not finalized.
        response_tx
            .send(Ok(BidiWriteObjectResponse {
                write_status: Some(WriteStatus::Resource(Object {
                    name: "test-obj".into(),
                    size: 100,
                    finalize_time: None,
                    ..Default::default()
                })),
                ..Default::default()
            }))
            .await?;
        tokio::task::yield_now().await;

        // Assert.
        // The caller is still waiting, rather than being told the upload finalized.
        assert!(matches!(
            finalize_rx.try_recv(),
            Err(oneshot::error::TryRecvError::Empty)
        ));

        // A genuinely finalized resource completes it.
        response_tx.send(Ok(finalized_response(100))).await?;
        let got = finalize_rx.await??;
        assert!(is_finalized(&got), "{got:?}");

        drop(tx);
        tokio::task::yield_now().await;
        drop(response_tx);
        handle.await??;
        Ok(())
    }

    /// Spawns a worker whose first stream is already connected and whose reconnect attempt yields
    /// a second stream, returning the handles needed to drive both.
    fn spawn_reconnecting_worker() -> ReconnectingWorkerContext {
        spawn_reconnecting_worker_with_capacity(DEFAULT_REPLAY_BUFFER_SIZE)
    }

    fn spawn_reconnecting_worker_with_capacity(capacity: usize) -> ReconnectingWorkerContext {
        let (captured_stream2_req_tx, captured_stream2_req_rx) =
            mpsc::channel::<mpsc::Receiver<BidiWriteObjectRequest>>(1);
        let (stream2_resp_tx, stream2_resp_rx) = mpsc::channel(10);
        let stream2 = TonicResponse::from(stream2_resp_rx);

        let mut mock = MockTestClient::new();
        mock.expect_start()
            .times(1)
            .return_once(move |_, _, req_rx, _, _, _| {
                let _ = captured_stream2_req_tx.try_send(req_rx);
                Ok(Ok(stream2))
            });

        let (handle, intent_tx, stream1_rx, stream1_resp_tx) =
            spawn_worker(mock, capacity, Some(test_spec_state()));
        (
            handle,
            intent_tx,
            stream1_rx,
            stream1_resp_tx,
            captured_stream2_req_rx,
            stream2_resp_tx,
        )
    }

    #[tokio::test]
    async fn run_reconnect_and_replay_unpersisted_chunks() -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        // Append two 10-byte chunks
        let chunk1 = bytes::Bytes::from_static(b"0123456789");
        let req1 = BidiWriteObjectRequest {
            write_offset: 0,
            data: Some(Data::ChecksummedData(ChecksummedData {
                content: chunk1.clone(),
                crc32c: Some(crc32c::crc32c(&chunk1)),
            })),
            ..Default::default()
        };
        let chunk2 = bytes::Bytes::from_static(b"abcdefghij");
        let req2 = BidiWriteObjectRequest {
            write_offset: 10,
            data: Some(Data::ChecksummedData(ChecksummedData {
                content: chunk2.clone(),
                crc32c: Some(crc32c::crc32c(&chunk2)),
            })),
            ..Default::default()
        };
        intent_tx.send(UploadIntent::Append(req1)).await?;
        intent_tx.send(UploadIntent::Append(req2)).await?;

        // Ensure both chunks were dispatched on stream 1 (and therefore buffered for replay).
        let s1_req1 = stream1_rx.recv().await.unwrap();
        assert_eq!(s1_req1.write_offset, 0);
        let s1_req2 = stream1_rx.recv().await.unwrap();
        assert_eq!(s1_req2.write_offset, 10);

        // Act.
        // Simulate stream 1 failure by dropping response stream
        drop(stream1_resp_tx);

        // Connector reconnects to stream 2; server initial message reports persisted_size = 10
        // (chunk 1 persisted)
        stream2_resp_tx
            .send(Ok(persisted_size_response(10)))
            .await?;

        // Assert.
        // Verify that stream 2 received the reconnect opening handshake request
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        assert!(initial_req.first_message.is_some());

        // Verify that chunk 2 (unpersisted) is replayed over stream 2!
        let replayed_req = stream2_req_rx.recv().await.unwrap();
        assert_eq!(replayed_req.write_offset, 10);
        if let Some(Data::ChecksummedData(cd)) = replayed_req.data {
            assert_eq!(cd.content, chunk2);
            assert_eq!(cd.crc32c, Some(crc32c::crc32c(&chunk2)));
        } else {
            panic!("expected ChecksummedData");
        }

        // Chunk 1 was acknowledged by the handshake, so it is not replayed, and nothing else
        // follows (no pending flush/finalize, and the buffer is far below the watermark).
        assert!(matches!(
            stream2_req_rx.try_recv(),
            Err(mpsc::error::TryRecvError::Empty)
        ));

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_reconnect_uses_refreshed_write_handle() -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        intent_tx.send(append_intent(0, 16)).await?;
        let _ = stream1_rx.recv().await.unwrap();

        // Act.
        // The service refreshes the write handle in a response that carries no `write_status`.
        let refreshed = BidiWriteHandle {
            handle: bytes::Bytes::from_static(b"refreshed-handle"),
        };
        stream1_resp_tx
            .send(Ok(BidiWriteObjectResponse {
                write_handle: Some(refreshed.clone()),
                ..Default::default()
            }))
            .await?;
        tokio::task::yield_now().await;

        // Stream 1 breaks; the reconnect handshake reports nothing persisted.
        drop(stream1_resp_tx);
        stream2_resp_tx.send(Ok(persisted_size_response(0))).await?;

        // Assert.
        // `test_spec_state()` starts with no write handle, so the handshake carries one only if
        // the worker recorded the refresh.
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        let Some(FirstMessage::AppendObjectSpec(spec)) = initial_req.first_message else {
            panic!("expected AppendObjectSpec handshake, got {initial_req:?}");
        };
        assert_eq!(spec.write_handle, Some(refreshed));

        // The unacknowledged chunk is still replayed after the handshake.
        let replayed = stream2_req_rx.recv().await.unwrap();
        assert_eq!(replayed.write_offset, 0);
        assert!(replayed.data.is_some(), "{replayed:?}");

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_stop_on_closed_requests() -> anyhow::Result<()> {
        let (handle, tx, _request_rx, _response_tx) = spawn_test_worker();
        drop(tx);
        tokio::task::yield_now().await;
        drop(_response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_handles_trailing_response_after_intents_close() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, _request_rx, response_tx) = spawn_test_worker();

        // Act.
        // Close the intent channel first so the worker drains the stream in
        // `wait_for_server_completion`, then deliver one late response.
        drop(tx);
        tokio::task::yield_now().await;
        response_tx.send(Ok(persisted_size_response(10))).await?;
        drop(response_tx);

        // Assert.
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_fails_on_stream_error_after_intents_close() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, _request_rx, response_tx) = spawn_test_worker();

        // Act.
        // Once the intent channel is closed the worker no longer reconnects, so a stream error
        // while draining is terminal.
        drop(tx);
        tokio::task::yield_now().await;
        response_tx
            .send(Err(Status::unavailable("try again")))
            .await?;

        // Assert.
        let err = handle.await?.unwrap_err().to_string();
        assert!(err.contains("try again"), "{err}");
        Ok(())
    }

    /// Spawns a worker whose reconnect attempt fails with `err`.
    fn spawn_failing_reconnect_worker(err: Error) -> TestWorkerContext {
        spawn_failing_reconnect_worker_with_capacity(err, DEFAULT_REPLAY_BUFFER_SIZE)
    }

    fn spawn_failing_reconnect_worker_with_capacity(
        err: Error,
        capacity: usize,
    ) -> TestWorkerContext {
        let mut mock = MockTestClient::new();
        mock.expect_start()
            .return_once(move |_, _, _, _, _, _| Err(err));
        spawn_worker(mock, capacity, Some(test_spec_state()))
    }

    #[tokio::test]
    async fn run_server_closes_unexpectedly() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, _request_rx, response_tx) =
            spawn_failing_reconnect_worker(permanent_error());

        // Act.
        // Close the stream from the server side unexpectedly while upload is not finalized.
        drop(response_tx);

        // Assert.
        let result = handle.await?;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("cannot serialize the request"));
        assert!(err.contains("PERMISSION_DENIED"));

        drop(tx);
        Ok(())
    }

    #[tokio::test]
    async fn run_stream_error_during_flush() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) =
            spawn_failing_reconnect_worker(permanent_error());

        let (flush_tx, flush_rx) = oneshot::channel();
        let flush_request = BidiWriteObjectRequest {
            flush: true,
            state_lookup: true,
            write_offset: 100,
            ..Default::default()
        };

        // Act.
        tx.send(UploadIntent::Flush(flush_request.clone(), flush_tx))
            .await?;

        let stream_req = request_rx.recv().await.unwrap();
        assert!(stream_req.flush);

        // Drop response stream and simulate failed reconnect
        drop(response_tx);

        // Assert.
        let received_resp = flush_rx.await?;
        assert!(received_resp.is_err());
        let err = received_resp.unwrap_err().to_string();
        assert!(err.contains("cannot serialize the request"));
        assert!(err.contains("PERMISSION_DENIED"));

        let result = handle.await?;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("cannot serialize the request"));
        assert!(err.contains("PERMISSION_DENIED"));
        Ok(())
    }

    #[tokio::test]
    async fn run_stream_error_then_queue_requests() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, _request_rx, response_tx) =
            spawn_failing_reconnect_worker(permanent_error());

        let (flush_tx1, flush_rx1) = oneshot::channel();
        let (flush_tx2, flush_rx2) = oneshot::channel();

        // Act.
        // Drop the server response stream to simulate the remote network crash.
        // The worker will wake up and attempt reconnect, which fails and triggers draining.
        drop(response_tx);

        // Put requests into the channel immediately. Because it has capacity 10
        // and we haven't yielded, these are queued in the requests buffer synchronously.
        let valid_flush = || BidiWriteObjectRequest {
            flush: true,
            state_lookup: true,
            write_offset: 100,
            ..Default::default()
        };
        tx.send(UploadIntent::Flush(valid_flush(), flush_tx1))
            .await?;
        tx.send(UploadIntent::Flush(valid_flush(), flush_tx2))
            .await?;

        // Assert.
        let payload1 = flush_rx1.await.unwrap();
        assert!(payload1.is_err());
        let err1 = payload1.unwrap_err().to_string();
        assert!(err1.contains("cannot serialize the request"));
        assert!(err1.contains("PERMISSION_DENIED"));

        let payload2 = flush_rx2.await.unwrap();
        assert!(payload2.is_err());
        let err2 = payload2.unwrap_err().to_string();
        assert!(err2.contains("cannot serialize the request"));
        assert!(err2.contains("PERMISSION_DENIED"));

        let result = handle.await?;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("cannot serialize the request"));
        assert!(err.contains("PERMISSION_DENIED"));

        Ok(())
    }

    #[tokio::test]
    async fn run_error_drains_intents_queued_behind_full_buffer() -> anyhow::Result<()> {
        // Arrange.
        // A single append fills the buffer, which disables the intent branch of the `run()` select.
        // Intents sent afterwards stay in the channel and can only be completed by the
        // channel-drain loop in `drain_intents_on_error`.
        let capacity = MIN_REPLAY_BUFFER_SIZE;
        let (handle, tx, mut request_rx, response_tx) =
            spawn_failing_reconnect_worker_with_capacity(permanent_error(), capacity);

        tx.send(append_intent(0, capacity)).await?;
        let stream_req = request_rx.recv().await.unwrap();
        // The append crossed the watermark, so the worker piggybacked a state lookup and is now
        // waiting for the response with a full buffer.
        assert!(stream_req.state_lookup);

        let (flush, flush_rx) = flush_intent(capacity as i64);
        tx.send(flush).await?;
        let (finalize, finalize_rx) = finalize_intent(capacity as i64);
        tx.send(finalize).await?;

        // Act.
        // Fail the stream. Reconnecting fails permanently, so the worker drains.
        drop(response_tx);

        // Assert.
        for received in [flush_rx.await?, finalize_rx.await?] {
            let err = received.unwrap_err().to_string();
            assert!(err.contains("PERMISSION_DENIED"), "{err}");
        }
        let err = handle.await?.unwrap_err().to_string();
        assert!(err.contains("PERMISSION_DENIED"), "{err}");
        Ok(())
    }

    #[tokio::test]
    async fn run_full_buffer_recovers_and_resumes_after_server_ack() -> anyhow::Result<()> {
        // Arrange.
        // Fill the replay buffer past `MIN_REPLAY_BUFFER_SIZE` (4 MiB + 1 byte) with three 2 MiB
        // appends so `is_full()` becomes true and the `run()` `tokio::select!` guard disables the
        // `requests.recv()` branch. Queue a fourth append while the buffer is full, then send a
        // server ack that drains the buffer and verify the fourth append is picked up and sent.
        let (handle, tx, mut request_rx, response_tx) =
            spawn_test_worker_with_capacity(MIN_REPLAY_BUFFER_SIZE);

        let chunk_len = MAX_WRITE_CHUNK_SIZE;
        for i in 0..3_i64 {
            tx.send(append_intent(i * chunk_len as i64, chunk_len))
                .await?;
            let dispatched = request_rx.recv().await.expect("chunk must be dispatched");
            assert_eq!(dispatched.write_offset, i * chunk_len as i64);
            // With `MIN_REPLAY_BUFFER_SIZE` the watermark is 1 byte, so the first append crosses
            // it and carries the `flush + state_lookup` probe. Later appends do not repeat it
            // while that probe is outstanding, which is what guarantees a server response is
            // owed by the time the buffer is full.
            let is_first = i == 0;
            assert_eq!(dispatched.flush, is_first, "{dispatched:?}");
            assert_eq!(dispatched.state_lookup, is_first, "{dispatched:?}");
        }

        // Queue a fourth append while the replay buffer is full (6 MiB >= 4 MiB + 1). Because the
        // intent branch is disabled, the worker must not forward it yet.
        let fourth_offset = 3 * chunk_len as i64;
        tx.send(append_intent(fourth_offset, TEST_CHUNK_SIZE))
            .await?;
        tokio::task::yield_now().await;
        assert!(
            request_rx.try_recv().is_err(),
            "fourth append must remain gated in the intent channel while replay_buffer.is_full()"
        );

        // Act.
        // Acknowledge all 6 MiB so `replay_buffer.ack()` drains the buffer and re-enables the
        // intent branch in the next loop iteration.
        response_tx
            .send(Ok(persisted_size_response(fourth_offset)))
            .await?;

        // Assert.
        let fourth = request_rx
            .recv()
            .await
            .expect("fourth append must be dispatched once the buffer drains");
        assert_eq!(fourth.write_offset, fourth_offset);

        drop(tx);
        tokio::task::yield_now().await;
        drop(response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_panic_on_flush_missing_state_lookup() {
        let (handle, tx, _request_rx, _response_tx) = spawn_test_worker();
        let (flush_tx, _flush_rx) = oneshot::channel();
        let flush_request = BidiWriteObjectRequest {
            flush: true,
            state_lookup: false, // Invalid
            ..Default::default()
        };
        let _ = tx.send(UploadIntent::Flush(flush_request, flush_tx)).await;
        assert!(handle.await.unwrap_err().is_panic());
    }

    #[tokio::test]
    async fn run_panic_on_flush_missing_flush() {
        let (handle, tx, _request_rx, _response_tx) = spawn_test_worker();
        let (flush_tx, _flush_rx) = oneshot::channel();
        let flush_request = BidiWriteObjectRequest {
            flush: false, // Invalid
            state_lookup: true,
            ..Default::default()
        };
        let _ = tx.send(UploadIntent::Flush(flush_request, flush_tx)).await;
        assert!(handle.await.unwrap_err().is_panic());
    }

    #[tokio::test]
    async fn run_panic_on_finalize_missing_finish_write() {
        let (handle, tx, _request_rx, _response_tx) = spawn_test_worker();
        let (finalize_tx, _finalize_rx) = oneshot::channel();
        let finalize_request = BidiWriteObjectRequest {
            finish_write: false, // Invalid
            flush: true,
            ..Default::default()
        };
        let _ = tx
            .send(UploadIntent::Finalize(finalize_request, finalize_tx))
            .await;
        assert!(handle.await.unwrap_err().is_panic());
    }

    #[tokio::test]
    async fn run_panic_on_finalize_missing_flush() {
        let (handle, tx, _request_rx, _response_tx) = spawn_test_worker();
        let (finalize_tx, _finalize_rx) = oneshot::channel();
        let finalize_request = BidiWriteObjectRequest {
            finish_write: true,
            flush: false, // Invalid
            ..Default::default()
        };
        let _ = tx
            .send(UploadIntent::Finalize(finalize_request, finalize_tx))
            .await;
        assert!(handle.await.unwrap_err().is_panic());
    }

    #[tokio::test]
    async fn run_append_injects_watermark_flush() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) =
            spawn_test_worker_with_capacity(TEST_CAPACITY);

        // Act.
        tx.send(append_intent(0, TEST_CHUNK_SIZE)).await?;
        tx.send(append_intent(TEST_CHUNK_SIZE as i64, TEST_CHUNK_SIZE))
            .await?;

        // Assert.
        // The first append stays below the watermark and is dispatched untouched.
        let first = request_rx.recv().await.unwrap();
        assert!(!first.flush, "{first:?}");
        assert!(!first.state_lookup, "{first:?}");

        // The second append crosses the watermark, so the worker piggybacks a flush and a
        // state_lookup on it. Only state_lookup elicits a server response.
        let second = request_rx.recv().await.unwrap();
        assert!(second.flush, "{second:?}");
        assert!(second.state_lookup, "{second:?}");

        drop(tx);
        tokio::task::yield_now().await;
        drop(response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_append_does_not_repeat_watermark_flush() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) =
            spawn_test_worker_with_capacity(TEST_CAPACITY);

        // Act.
        // Three appends, all above the watermark from the second one onwards, with no server
        // response in between.
        for i in 0..3 {
            tx.send(append_intent((i * TEST_CHUNK_SIZE) as i64, TEST_CHUNK_SIZE))
                .await?;
        }

        // Assert.
        let _first = request_rx.recv().await.unwrap();
        let second = request_rx.recv().await.unwrap();
        assert!(second.state_lookup, "{second:?}");

        // A state_lookup is already outstanding, so the third append is not flagged.
        let third = request_rx.recv().await.unwrap();
        assert!(!third.flush, "{third:?}");
        assert!(!third.state_lookup, "{third:?}");

        drop(tx);
        tokio::task::yield_now().await;
        drop(response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_appends_past_capacity_without_explicit_flush() -> anyhow::Result<()> {
        // Arrange.
        // Send more than twice `TEST_CAPACITY` without ever flushing. The replay buffer cannot
        // hold it all, so the worker only finishes if its watermark probes keep eliciting the
        // server responses that drive `ack()`. Otherwise `is_full()` disables the intent branch,
        // the worker never observes the closed channel, and the timeout below fires.
        const APPEND_COUNT: usize = 6;
        const { assert!(APPEND_COUNT * MAX_WRITE_CHUNK_SIZE > 2 * TEST_CAPACITY) }
        let (handle, tx, mut request_rx, response_tx) =
            spawn_test_worker_with_capacity(TEST_CAPACITY);

        // A server that replies only when the client asks for it via state_lookup.
        let server = tokio::spawn(async move {
            let mut data_requests = 0_usize;
            let mut persisted_size = 0_i64;
            while let Some(request) = request_rx.recv().await {
                if let Some(Data::ChecksummedData(cd)) = request.data.as_ref() {
                    data_requests += 1;
                    persisted_size = request.write_offset + cd.content.len() as i64;
                }
                if request.state_lookup
                    && response_tx
                        .send(Ok(persisted_size_response(persisted_size)))
                        .await
                        .is_err()
                {
                    break;
                }
            }
            data_requests
        });

        // Act.
        let result = tokio::time::timeout(std::time::Duration::from_secs(10), async move {
            for i in 0..APPEND_COUNT {
                let write_offset = (i * MAX_WRITE_CHUNK_SIZE) as i64;
                tx.send(append_intent(write_offset, MAX_WRITE_CHUNK_SIZE))
                    .await
                    .expect("the worker must keep accepting appends");
            }
            drop(tx);
            handle.await.expect("the worker task must not panic")
        })
        .await
        .expect("appends without an explicit flush must not deadlock");

        // Assert.
        result?;
        assert_eq!(server.await?, APPEND_COUNT);
        Ok(())
    }

    #[tokio::test]
    async fn run_reconnect_restores_watermark_state_lookup() -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker_with_capacity(TEST_CAPACITY);

        // Fill the replay buffer up to the watermark on the first stream.
        intent_tx.send(append_intent(0, TEST_CHUNK_SIZE)).await?;
        intent_tx
            .send(append_intent(TEST_CHUNK_SIZE as i64, TEST_CHUNK_SIZE))
            .await?;
        let _ = stream1_rx.recv().await.unwrap();
        let _ = stream1_rx.recv().await.unwrap();

        // Act.
        // Break the first stream. The watermark state_lookup sent on it is lost.
        drop(stream1_resp_tx);

        // The reconnected stream reports nothing persisted, so both chunks are replayed.
        stream2_resp_tx.send(Ok(persisted_size_response(0))).await?;

        // Assert.
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        assert!(initial_req.first_message.is_some());

        // Replayed chunks carry data but neither flush nor state_lookup.
        for i in 0..2 {
            let replayed = stream2_req_rx.recv().await.unwrap();
            assert_eq!(replayed.write_offset, (i * TEST_CHUNK_SIZE) as i64);
            assert!(replayed.data.is_some(), "{replayed:?}");
            assert!(!replayed.flush, "{replayed:?}");
            assert!(!replayed.state_lookup, "{replayed:?}");
        }

        // The buffer is still at the watermark and no user request is pending, so the worker
        // re-establishes the invariant with a standalone state_lookup.
        let watermark_req = stream2_req_rx.recv().await.unwrap();
        assert!(watermark_req.flush, "{watermark_req:?}");
        assert!(watermark_req.state_lookup, "{watermark_req:?}");
        assert_eq!(watermark_req.write_offset, (2 * TEST_CHUNK_SIZE) as i64);

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    fn finalize_intent(
        write_offset: i64,
    ) -> (
        UploadIntent,
        oneshot::Receiver<crate::Result<BidiWriteObjectResponse>>,
    ) {
        let (finalize_tx, finalize_rx) = oneshot::channel();
        let intent = UploadIntent::Finalize(
            BidiWriteObjectRequest {
                write_offset,
                flush: true,
                finish_write: true,
                ..Default::default()
            },
            finalize_tx,
        );
        (intent, finalize_rx)
    }

    #[tokio::test]
    async fn run_reconnect_resends_pending_finalize() -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        let (intent, mut finalize_rx) = finalize_intent(50);
        intent_tx.send(intent).await?;
        let sent_finalize = stream1_rx.recv().await.unwrap();
        assert!(sent_finalize.finish_write);

        // Act.
        // Stream 1 breaks before delivering the finalize response.
        drop(stream1_resp_tx);

        // The reconnect handshake returns a resource without `finalize_time`, which is what a
        // takeover stream reports. It is not a finalization.
        stream2_resp_tx
            .send(Ok(BidiWriteObjectResponse {
                write_status: Some(WriteStatus::Resource(Object {
                    name: "test-object".into(),
                    size: 50,
                    finalize_time: None,
                    ..Default::default()
                })),
                ..Default::default()
            }))
            .await?;

        // Assert.
        // The finalize is re-sent on the new stream after the handshake, ...
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        assert!(initial_req.first_message.is_some());
        let resent = stream2_req_rx.recv().await.unwrap();
        assert!(resent.finish_write, "{resent:?}");
        assert_eq!(resent.write_offset, 50);

        // ... and the caller has not been told the upload finalized. `resent` was sent after the
        // worker processed the handshake, so the handshake has already had its chance to (wrongly)
        // complete the pending finalize.
        assert!(matches!(
            finalize_rx.try_recv(),
            Err(oneshot::error::TryRecvError::Empty)
        ));

        // The service then answers the re-sent finalize.
        stream2_resp_tx.send(Ok(finalized_response(50))).await?;
        let got = finalize_rx.await??;
        assert!(is_finalized(&got), "{got:?}");

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_reconnect_completes_already_finalized_request_without_resending()
    -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        let (intent, finalize_rx) = finalize_intent(50);
        intent_tx.send(intent).await?;
        let sent_finalize = stream1_rx.recv().await.unwrap();
        assert!(sent_finalize.finish_write);

        // Act.
        // Stream 1 breaks after the service finalized the object but before it could deliver the
        // response.
        drop(stream1_resp_tx);
        let finalized = finalized_response(50);
        stream2_resp_tx.send(Ok(finalized.clone())).await?;

        // Assert.
        // The handshake response completes the pending finalize.
        let got_resp = finalize_rx.await??;
        assert_eq!(got_resp, finalized);

        // Stream 2 only received the initial AppendObjectSpec handshake, NOT a duplicate Finalize.
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        assert!(initial_req.first_message.is_some());
        assert!(stream2_req_rx.try_recv().is_err());

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_reconnect_replays_chunk_and_resends_pending_flush() -> anyhow::Result<()> {
        // Arrange.
        // `transport.flush()` queues the residual append immediately followed by the flush, so
        // `[Append, Flush]` is the usual shape of the worker's state when a stream breaks.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        intent_tx.send(append_intent(0, 10)).await?;
        let (intent, mut flush_rx) = flush_intent(10);
        intent_tx.send(intent).await?;
        let _ = stream1_rx.recv().await.unwrap();
        let sent_flush = stream1_rx.recv().await.unwrap();
        assert!(sent_flush.flush);

        // Act.
        // Stream 1 breaks before the append reached the service; the handshake reports nothing
        // persisted.
        drop(stream1_resp_tx);
        stream2_resp_tx.send(Ok(persisted_size_response(0))).await?;

        // Assert.
        // Stream 2 receives the handshake, then the replayed chunk, then the re-sent flush.
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        assert!(initial_req.first_message.is_some());

        let replayed = stream2_req_rx.recv().await.unwrap();
        assert_eq!(replayed.write_offset, 0);
        assert!(replayed.data.is_some(), "{replayed:?}");
        assert!(!replayed.flush, "{replayed:?}");
        assert!(!replayed.state_lookup, "{replayed:?}");

        let resent_flush = stream2_req_rx.recv().await.unwrap();
        assert_eq!(resent_flush.write_offset, 10);
        assert!(resent_flush.flush, "{resent_flush:?}");
        assert!(resent_flush.state_lookup, "{resent_flush:?}");
        assert!(resent_flush.data.is_none(), "{resent_flush:?}");

        // The caller is still waiting: the handshake did not cover the flush offset.
        assert!(matches!(
            flush_rx.try_recv(),
            Err(oneshot::error::TryRecvError::Empty)
        ));

        // The service answers the re-sent flush.
        let server_resp = persisted_size_response(10);
        stream2_resp_tx.send(Ok(server_resp.clone())).await?;
        let got = flush_rx.await??;
        assert_eq!(got.write_status, server_resp.write_status);

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_reconnect_handshake_completes_pending_flush_without_resending()
    -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        intent_tx.send(append_intent(0, 10)).await?;
        let (intent, flush_rx) = flush_intent(10);
        intent_tx.send(intent).await?;
        let _ = stream1_rx.recv().await.unwrap();
        let _ = stream1_rx.recv().await.unwrap();

        // Act.
        // Stream 1 breaks after the service persisted the append but before the flush response
        // was delivered; the handshake already reports the flush offset as persisted.
        drop(stream1_resp_tx);
        let handshake = persisted_size_response(10);
        stream2_resp_tx.send(Ok(handshake.clone())).await?;

        // Assert.
        // The handshake response completes the pending flush, ...
        let got = flush_rx.await??;
        assert_eq!(got.write_status, handshake.write_status);

        // ... and stream 2 receives only the handshake: no replayed chunk, no duplicate flush.
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        assert!(initial_req.first_message.is_some());
        assert!(matches!(
            stream2_req_rx.try_recv(),
            Err(mpsc::error::TryRecvError::Empty)
        ));

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_send_error_defers_reconnect_to_rx_and_replays_chunk() -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            stream1_rx,
            stream1_resp_tx,
            mut captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        // Prepare Stream 2's initial handshake response (`PersistedSize(0)`).
        stream2_resp_tx.send(Ok(persisted_size_response(0))).await?;

        // Act.
        // Drop Stream 1's request receiver so `tx.send(request).await` fails after
        // `process_intent` buffers the chunk, then deliver a transient gRPC status on Stream 1's
        // response half so `rx.next_message()` drives `reconnect`.
        drop(stream1_rx);
        intent_tx.send(append_intent(0, 16)).await?;
        tokio::task::yield_now().await;
        stream1_resp_tx
            .send(Err(Status::unavailable("stream broken")))
            .await?;

        // Assert.
        // `reconnect` replaces the stream halves and replays the chunk that `process_intent`
        // buffered before `tx.send` failed.
        let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap();
        let initial_req = stream2_req_rx.recv().await.unwrap();
        assert!(initial_req.first_message.is_some());

        let replayed_req = stream2_req_rx.recv().await.unwrap();
        assert_eq!(replayed_req.write_offset, 0);
        let Some(Data::ChecksummedData(cd)) = replayed_req.data else {
            panic!("expected ChecksummedData on replayed append");
        };
        assert_eq!(cd.content.len(), 16);

        drop(intent_tx);
        tokio::task::yield_now().await;
        drop(stream2_resp_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_watermark_rearms_after_non_draining_ack() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) =
            spawn_test_worker_with_capacity(TEST_CAPACITY);

        // Send two chunks to hit the watermark (`2 * TEST_CHUNK_SIZE`).
        tx.send(append_intent(0, TEST_CHUNK_SIZE)).await?;
        tx.send(append_intent(TEST_CHUNK_SIZE as i64, TEST_CHUNK_SIZE))
            .await?;
        let _first = request_rx.recv().await.unwrap();
        let second = request_rx.recv().await.unwrap();
        assert!(second.state_lookup);

        // Act.
        // Server responds with PersistedSize(0), clearing self_flush_outstanding without draining
        // the buffer.
        response_tx.send(Ok(persisted_size_response(0))).await?;

        // Assert.
        // Worker immediately dispatches a standalone watermark state_lookup at offset 2 *
        // TEST_CHUNK_SIZE.
        let rearmed = request_rx.recv().await.unwrap();
        assert!(rearmed.flush, "{rearmed:?}");
        assert!(rearmed.state_lookup, "{rearmed:?}");
        assert_eq!(rearmed.write_offset, (2 * TEST_CHUNK_SIZE) as i64);

        drop(tx);
        tokio::task::yield_now().await;
        drop(response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_response_missing_write_status_does_not_trip_regression_or_flush()
    -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) = spawn_test_worker();

        tx.send(append_intent(0, 10)).await?;
        let _ = request_rx.recv().await.unwrap();

        // Acknowledge 10 bytes first so worker.persisted_size = 10.
        response_tx.send(Ok(persisted_size_response(10))).await?;
        tokio::task::yield_now().await;

        let (intent, mut flush_rx) = flush_intent(10);
        tx.send(intent).await?;
        let _ = request_rx.recv().await.unwrap();

        // Act.
        // Send a response with write_status: None (e.g. write_handle refresh only). It must not be
        // treated as persisted_size = 0 (which would fail the regression check) and must not
        // satisfy the pending flush.
        response_tx
            .send(Ok(BidiWriteObjectResponse {
                write_status: None,
                ..Default::default()
            }))
            .await?;
        tokio::task::yield_now().await;

        // Assert.
        assert!(matches!(
            flush_rx.try_recv(),
            Err(oneshot::error::TryRecvError::Empty)
        ));

        response_tx.send(Ok(persisted_size_response(10))).await?;
        flush_rx.await??;

        drop(tx);
        tokio::task::yield_now().await;
        drop(response_tx);
        handle.await??;
        Ok(())
    }

    #[tokio::test]
    async fn run_reconnect_fails_on_persisted_size_regression() -> anyhow::Result<()> {
        // Arrange.
        let (
            handle,
            intent_tx,
            mut stream1_rx,
            stream1_resp_tx,
            _captured_stream2_req_rx,
            stream2_resp_tx,
        ) = spawn_reconnecting_worker();

        intent_tx.send(append_intent(0, 20)).await?;
        let _ = stream1_rx.recv().await.unwrap();

        // Server acknowledges 15 bytes, trimming [0..15) from the replay buffer.
        stream1_resp_tx
            .send(Ok(persisted_size_response(15)))
            .await?;
        tokio::task::yield_now().await;

        // Act.
        // Stream 1 breaks; Stream 2 claims only 10 bytes persisted (< 15 already trimmed).
        drop(stream1_resp_tx);
        stream2_resp_tx
            .send(Ok(persisted_size_response(10)))
            .await?;

        // Assert.
        let err = handle.await?.unwrap_err();
        assert!(
            err.to_string()
                .contains("regressed below acknowledged offset"),
            "{err:?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn run_stream_permanent_grpc_error_fails_without_reconnecting() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, _request_rx, response_tx) = spawn_test_worker();

        // Act.
        // Send a permanent gRPC error (INVALID_ARGUMENT) on the stream. `Worker::reconnect` is
        // entered, but `Connector::reconnect` rejects the permanent error via the retry policy
        // before opening a new stream, so `start()` is never called; `spawn_test_worker` sets
        // `mock.expect_start().never()` to prove it.
        response_tx
            .send(Err(Status::invalid_argument("bad write_offset")))
            .await?;

        // Assert.
        let err = handle.await?.unwrap_err();
        assert!(err.to_string().contains("bad write_offset"), "{err:?}");
        drop(tx);
        Ok(())
    }

    #[tokio::test]
    async fn run_fails_when_server_reports_unexpected_finalized_object() -> anyhow::Result<()> {
        // Arrange.
        let (handle, tx, mut request_rx, response_tx) = spawn_test_worker();

        tx.send(append_intent(0, 10)).await?;
        let _ = request_rx.recv().await.unwrap();
        drop(tx);
        tokio::task::yield_now().await;

        // Act.
        // Server returns a finalized Object resource while draining in
        // `wait_for_server_completion` even though no Finalize intent was sent.
        response_tx.send(Ok(finalized_response(10))).await?;

        // Assert.
        let err = handle.await?.unwrap_err();
        assert!(
            err.to_string().contains("object is already finalized"),
            "{err:?}"
        );
        Ok(())
    }
}