diff --git a/src/main/java/com/google/genai/gaos/AsyncInteractions.java b/src/main/java/com/google/genai/gaos/AsyncInteractions.java index 568840d1860..fa487f5b96c 100644 --- a/src/main/java/com/google/genai/gaos/AsyncInteractions.java +++ b/src/main/java/com/google/genai/gaos/AsyncInteractions.java @@ -31,13 +31,11 @@ import com.google.genai.gaos.models.operations.async.CancelInteractionByIdRequestBuilder; import com.google.genai.gaos.models.operations.async.CancelInteractionByIdResponse; import com.google.genai.gaos.models.operations.async.CreateInteractionRequestBuilder; -import com.google.genai.gaos.models.operations.async.CreateInteractionResponse; import com.google.genai.gaos.models.operations.async.DeleteInteractionRequestBuilder; import com.google.genai.gaos.models.operations.async.DeleteInteractionResponse; import com.google.genai.gaos.models.operations.async.GetInteractionByIdRequestBuilder; import com.google.genai.gaos.models.operations.async.GetInteractionByIdResponse; import com.google.genai.gaos.operations.CancelInteractionById; -import com.google.genai.gaos.operations.CreateInteraction; import com.google.genai.gaos.operations.DeleteInteraction; import com.google.genai.gaos.operations.GetInteractionById; import com.google.genai.gaos.operations.Operations; @@ -91,6 +89,20 @@ public CompletableFuture> create(@Nonnull return create(null, body, null); } + /** + * Creates a new interaction. + * + * @param body The request body. + * @param automaticContinuation Whether to automatically continue decoding across + * continuation hops; see {@link #create(CreateInteractionRequest, Options)}. + * @return {@code CompletableFuture}> - The async response + */ + public CompletableFuture> create( + @Nonnull CreateInteractionRequestBody body, boolean automaticContinuation) { + return create( + new CreateInteractionRequest(null, body, automaticContinuation, null), null); + } + /** * Creates a new interaction. * @@ -103,19 +115,42 @@ public CompletableFuture> create(@Nonnull public CompletableFuture> create( @Nullable String apiVersion, @Nonnull CreateInteractionRequestBody body, @Nullable Options options) { - CreateInteractionRequest request = new CreateInteractionRequest(apiVersion, body); - AsyncRequestOperation operation - = new CreateInteraction.Async( - sdkConfiguration, options, sdkConfiguration.retryScheduler(), - _headers); - return Operations.relayCancel(Operations.applyBodyReadAsync(operation.doRequest(request), - operation::handleResponse) - .thenApplyAsync(response -> new EventStream( - response.rawResponse().body(), - new TypeReference() { - }, - Utils.mapper(), - Optional.of("[DONE]")), Operations.streamCompletionExecutor()), operation); + return create(new CreateInteractionRequest(apiVersion, body), options); + } + + /** + * Creates a new interaction. + * + * @param request The request object containing all the parameters for the API call. + * @return A CompletableFuture that completes with a blocking event stream once response headers are received. + * Iterating the stream blocks the calling thread; close it after use. + */ + public CompletableFuture> create( + @Nonnull CreateInteractionRequest request) { + return create(request, null); + } + + /** + * Creates a new interaction. + * + *

Automatic continuation ({@code request.automaticContinuation()}, default {@code true}) + * works as described on {@link Interactions#create(CreateInteractionRequest, Options)}: with + * {@code stream=true}, the events of every continuation hop are yielded sequentially and + * verbatim on the returned stream, which fetches the next hop on the thread iterating it. A + * failure after progress ends the iteration with a + * {@link com.google.genai.gaos.models.errors.ContinuationInterruptedException} carrying the + * token to resume from. {@code background=true} with {@code automaticContinuation(false)} is + * rejected with an {@link IllegalArgumentException} before any request is sent. + * + * @param request The request object containing all the parameters for the API call. + * @param options additional options + * @return A CompletableFuture that completes with a blocking event stream once response headers are received. + * Iterating the stream blocks the calling thread; close it after use. + */ + public CompletableFuture> create( + @Nonnull CreateInteractionRequest request, + @Nullable Options options) { + return InteractionContinuation.createAsync(sdkConfiguration, request, options, _headers); } diff --git a/src/main/java/com/google/genai/gaos/ContinuationEventStream.java b/src/main/java/com/google/genai/gaos/ContinuationEventStream.java new file mode 100644 index 00000000000..2dc813459de --- /dev/null +++ b/src/main/java/com/google/genai/gaos/ContinuationEventStream.java @@ -0,0 +1,516 @@ +/* + * 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 + * + * http://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. + */ + +package com.google.genai.gaos; + +import com.google.genai.gaos.models.errors.ContinuationInterruptedException; +import com.google.genai.gaos.models.errors.GaosBaseException; +import com.google.genai.gaos.models.interactions.ErrorEvent; +import com.google.genai.gaos.models.interactions.InteractionCompletedEvent; +import com.google.genai.gaos.models.interactions.InteractionCreatedEvent; +import com.google.genai.gaos.models.interactions.InteractionSSEEvent; +import com.google.genai.gaos.models.interactions.InteractionSSEStreamEvent; +import com.google.genai.gaos.models.interactions.InteractionSseEventInteraction; +import com.google.genai.gaos.models.interactions.InteractionStatusUpdate; +import com.google.genai.gaos.utils.EventStream; +import jakarta.annotation.Nullable; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.Optional; + +/** + * An {@link EventStream} of a streaming {@code interactions.create()} call that automatically + * continues across hops (HTTP requests). + * + *

Every event of every hop is yielded verbatim and in order, including each hop's own {@code + * interaction.created} and intermediate {@code interaction.completed} (with status {@code + * requires_continuation}) events. When a hop completes with {@code requires_continuation} and a + * token, the next hop is requested from that token. When the transport drops (or the stream ends) + * mid-hop after an {@code interaction.status_update} checkpoint token, and nothing was yielded + * since that checkpoint, the hop is resumed from it; a hop that drops before making progress is + * retried from its start token up to {@link InteractionContinuation#MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS} + * times in a row. SSE {@code error} events are terminal; HTTP status errors are not retried here. + * + *

Hand-written (not generated). Per-hop state lives in a {@link Hop} record replaced wholesale + * when a hop starts, so per-hop fields cannot leak across hops; loop state (next resume token and + * owner, no-progress reconnect counter, pending transport error) persists on the stream. + */ +public final class ContinuationEventStream extends EventStream { + + /** Opens the stream of a continuation hop. */ + @FunctionalInterface + public interface HopOpener { + /** + * Sends a continuation request resuming from {@code token}. + * + * @param token the {@code continuation_token} to resume from + * @param tokenOwnerId the id of the interaction that issued {@code token}, or null for the + * caller's own token + * @return the hop's event stream, or null if the response was not a stream (ends the stream) + */ + @Nullable + EventStream open(String token, @Nullable String tokenOwnerId) + throws Exception; + } + + private static final String ERROR_STATUS = "error"; + private static final String IN_PROGRESS = "in_progress"; + private static final String REQUIRES_CONTINUATION = "requires_continuation"; + + /** What was observed during one hop (HTTP request) of a continuation stream. */ + private static final class Hop { + // The continuation_token and owner id this hop was requested with. + @Nullable final String startToken; + @Nullable final String startTokenOwnerId; + // Seeded with startToken so a disconnect before the hop's first checkpoint retries from the + // hop start. + @Nullable String token; + // The interaction that issued `token`. With store=true the server only accepts a token + // together with the interaction that issued it as previous_interaction_id, so a retry of the + // seed token keeps the seed's owner even after this hop's own interaction was created. + @Nullable String tokenOwnerId; + // This hop's own interaction id, once known. + @Nullable String interactionId; + // Latest status observed during the hop, or null if none arrived yet. + @Nullable String status; + // Whether the HTTP stream was opened (response headers received). + boolean streamOpened; + // Whether output events were yielded since the latest checkpoint token. + boolean contentSinceCheckpoint; + + Hop(@Nullable String startToken, @Nullable String startTokenOwnerId, boolean streamOpened) { + this.startToken = startToken; + this.startTokenOwnerId = startTokenOwnerId; + this.token = startToken; + this.tokenOwnerId = startTokenOwnerId; + this.streamOpened = streamOpened; + } + } + + @Nullable private final HopOpener opener; + private final InteractionContinuation.CallDeadline deadline; + private final boolean store; + private final boolean warnWhenIncomplete; + + @Nullable private EventStream current; + private boolean finished = false; + private boolean closed = false; + + // Loop state. + @Nullable private String nextToken; + // The caller's own token: a failure that leaves nothing newer to resume from is not "after + // progress" and is thrown as-is. + @Nullable private final String initialToken; + @Nullable private String nextTokenOwnerId = null; + private final List hopIds = new ArrayList<>(); + private int noProgressReconnects = 0; + // The latest transport error not followed by progress; thrown if the loop gives up so the + // failure is not hidden by later clean closes. + @Nullable private Exception lastError = null; + private Hop hop; + + /** + * Creates a stream that continues {@code initial} across hops. + * + * @param initial the first hop's stream (its HTTP response headers were already received) + * @param opener opens the stream of each further hop + * @param initialContinuationToken the caller's own {@code continuation_token}, if any + * @param deadline the whole-call deadline, checked before each hop and before each event + */ + public ContinuationEventStream( + EventStream initial, + HopOpener opener, + @Nullable String initialContinuationToken, + InteractionContinuation.CallDeadline deadline) { + this(initial, opener, initialContinuationToken, deadline, /* store= */ true, false); + } + + ContinuationEventStream( + EventStream initial, + HopOpener opener, + @Nullable String initialContinuationToken, + InteractionContinuation.CallDeadline deadline, + boolean store) { + this(initial, opener, initialContinuationToken, deadline, store, false); + } + + private ContinuationEventStream( + EventStream initial, + @Nullable HopOpener opener, + @Nullable String initialContinuationToken, + InteractionContinuation.CallDeadline deadline, + boolean store, + boolean warnWhenIncomplete) { + super(); + this.current = initial; + this.opener = opener; + this.deadline = deadline; + this.store = store; + this.warnWhenIncomplete = warnWhenIncomplete; + this.initialToken = initialContinuationToken; + this.nextToken = initialContinuationToken; + this.hop = new Hop(initialContinuationToken, null, /* streamOpened= */ true); + } + + /** + * Wraps a stream that is not continued ({@code automaticContinuation=false}) so that a one-time + * warning is logged if it stops at {@code requires_continuation} with a token. + */ + static ContinuationEventStream withIncompleteWarning( + EventStream stream) { + return new ContinuationEventStream( + stream, null, null, new InteractionContinuation.CallDeadline(null, false), false, true); + } + + // --------------------------------------------------------------------------------------------- + // EventStream + // --------------------------------------------------------------------------------------------- + + @Override + public Optional next() throws IOException { + while (!finished && !closed) { + Exception error = null; + if (current == null) { + // Start the next hop. + String token = nextToken; + String tokenOwnerId = nextTokenOwnerId; + hop = new Hop(token, tokenOwnerId, /* streamOpened= */ false); + try { + deadline.check(); + EventStream next = opener.open(token, tokenOwnerId); + if (next == null) { + finished = true; + return Optional.empty(); + } + hop.streamOpened = true; + current = next; + } catch (Exception e) { + error = e; + } + } + if (error == null) { + Optional event; + try { + deadline.check(); + event = current.next(); + } catch (Exception e) { + error = e; + event = Optional.empty(); + } + if (event.isPresent()) { + observe(event.get()); + return event; + } + } + // The hop ended: cleanly, or with `error`. + closeCurrentQuietly(); + if (opener == null) { + finished = true; + if (error != null) { + throw rethrow(error); + } + if (warnWhenIncomplete + && REQUIRES_CONTINUATION.equals(hop.status) + && hop.token != null) { + InteractionContinuation.warnIncompleteContinuation(); + } + return Optional.empty(); + } + if (!endHop(error)) { + finished = true; + return Optional.empty(); + } + } + return Optional.empty(); + } + + @Override + public void close() throws IOException { + closed = true; + finished = true; + super.close(); + if (current != null) { + EventStream stream = current; + current = null; + stream.close(); + } + } + + @Override + public boolean isClosed() { + return closed; + } + + private void closeCurrentQuietly() { + if (current == null) { + return; + } + EventStream stream = current; + current = null; + try { + stream.close(); + } catch (IOException | RuntimeException ignored) { + // The hop is over; a failure to close a dead stream is not interesting. + } + } + + // --------------------------------------------------------------------------------------------- + // State machine (port of the Python SDK's StreamContinuation) + // --------------------------------------------------------------------------------------------- + + /** Records one streamed SSE event. */ + private void observe(InteractionSSEStreamEvent streamEvent) { + hop.streamOpened = true; + InteractionSSEEvent event = streamEvent.data().orElse(null); + String eventInteractionId = eventInteractionId(event); + if (eventInteractionId != null) { + hop.interactionId = eventInteractionId; + if (!hopIds.contains(eventInteractionId)) { + hopIds.add(eventInteractionId); + } + } + String previousToken = hop.token; + extractHopContinuationState(event); + if (hop.token != null && !hop.token.equals(previousToken)) { + hop.tokenOwnerId = hop.interactionId; + } + hop.contentSinceCheckpoint = hasContentSinceCheckpoint(event, hop.contentSinceCheckpoint); + } + + /** Updates the hop's {@code (status, token)} from one SSE event. */ + private void extractHopContinuationState(@Nullable InteractionSSEEvent event) { + if (event instanceof ErrorEvent) { + hop.status = ERROR_STATUS; + hop.token = null; + return; + } + InteractionSseEventInteraction interaction = eventInteraction(event); + if (interaction != null) { + String status = interaction.status().map(s -> s.value()).orElse(null); + if (status != null && !status.isEmpty()) { + hop.status = status; + } + String token = nonEmpty(interaction.continuationToken().orElse(null)); + if (event instanceof InteractionCompletedEvent) { + // The terminal snapshot is authoritative: a missing token clears the hop's seed token and + // any earlier mid-hop checkpoint token. + hop.token = token; + } else if (token != null) { + hop.token = token; + } + } + if (event instanceof InteractionStatusUpdate) { + InteractionStatusUpdate update = (InteractionStatusUpdate) event; + String status = update.status().map(s -> s.value()).orElse(null); + if (status != null && !status.isEmpty()) { + hop.status = status; + } + String token = nonEmpty(update.continuationToken().orElse(null)); + if (token != null) { + hop.token = token; + } + } + } + + /** Tracks whether output was yielded since the latest resume point. */ + private static boolean hasContentSinceCheckpoint( + @Nullable InteractionSSEEvent event, boolean contentSinceCheckpoint) { + if (event instanceof InteractionStatusUpdate) { + if (nonEmpty(((InteractionStatusUpdate) event).continuationToken().orElse(null)) != null) { + return false; + } + return contentSinceCheckpoint; + } + if (event instanceof InteractionCompletedEvent) { + // The terminal snapshot's token (if any) is a fresh resume point. + return false; + } + if (event instanceof InteractionCreatedEvent) { + // Carries no model output. + return contentSinceCheckpoint; + } + return true; + } + + @Nullable + private static InteractionSseEventInteraction eventInteraction( + @Nullable InteractionSSEEvent event) { + if (event instanceof InteractionCreatedEvent) { + return ((InteractionCreatedEvent) event).interaction().orElse(null); + } + if (event instanceof InteractionCompletedEvent) { + return ((InteractionCompletedEvent) event).interaction().orElse(null); + } + return null; + } + + /** The id of the interaction (hop) an SSE event belongs to, if present. */ + @Nullable + private static String eventInteractionId(@Nullable InteractionSSEEvent event) { + if (event instanceof InteractionStatusUpdate) { + String id = nonEmpty(((InteractionStatusUpdate) event).interactionId().orElse(null)); + if (id != null) { + return id; + } + } + InteractionSseEventInteraction interaction = eventInteraction(event); + return interaction == null ? null : nonEmpty(interaction.id().orElse(null)); + } + + private static boolean canResumeStreamHop(@Nullable String status, @Nullable String token) { + return token != null + && (status == null || IN_PROGRESS.equals(status) || REQUIRES_CONTINUATION.equals(status)); + } + + private boolean canResumeHop() { + // Resuming from the checkpoint would re-generate (and duplicate) any output yielded since it. + return !hop.contentSinceCheckpoint && canResumeStreamHop(hop.status, hop.token); + } + + /** Updates the no-progress reconnect counter after a resumable hop ends. */ + private static int nextNoProgressReconnects( + int noProgressReconnects, + @Nullable String status, + @Nullable String startToken, + @Nullable String token) { + if (REQUIRES_CONTINUATION.equals(status)) { + // Clean hop boundary (per-request decode limit reached). + return 0; + } + if (!Objects.equals(token, startToken)) { + // Dropped mid-hop after a new checkpoint token was received. + return 0; + } + // Dropped mid-hop before any new checkpoint: retry the hop's own token. + return noProgressReconnects + 1; + } + + /** + * The error to throw when the stream gives up with {@code error}. + * + *

If the server issued a resumable token during this call (a completed {@code + * requires_continuation} hop, or a mid-hop checkpoint), the work so far is resumable and {@code + * error} is wrapped in {@link ContinuationInterruptedException} carrying that token. Otherwise + * {@code error} is returned as-is. + */ + private Exception fail(Exception error) { + String token = canResumeStreamHop(hop.status, hop.token) ? hop.token : null; + if (token == null || token.equals(initialToken)) { + return error; + } + return new ContinuationInterruptedException( + "Automatic continuation of the stream was interrupted: " + + error + + ". The continuation_token can be used to resume (output yielded after a mid-hop" + + " checkpoint is re-generated).", + error, + null, + token, + store ? hop.tokenOwnerId : null, + new ArrayList<>(hopIds)); + } + + /** + * Ends the current hop and returns true if another hop should start. + * + *

Throws {@code error} (or the pending prior disconnect error) when the stream cannot safely + * resume or exceeds {@link InteractionContinuation#MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS}, + * wrapped in {@link ContinuationInterruptedException} if a server-issued token to resume from + * exists (see {@link #fail}). + */ + private boolean endHop(@Nullable Exception error) throws IOException { + if (error instanceof InteractionContinuation.ContinuationTimeoutException) { + throw rethrow(fail(error)); + } + if (error != null && !hop.streamOpened) { + // Pre-stream HTTP status or validation errors (already retried by RetryConfig) must + // propagate immediately; only connection errors resume. + if (!isConnectionError(error)) { + throw rethrow(fail(error)); + } + } + + if (!canResumeHop()) { + if (error != null) { + throw rethrow(fail(error)); + } + return false; + } + + if (REQUIRES_CONTINUATION.equals(hop.status) && Objects.equals(hop.token, hop.startToken)) { + // The hop finished but handed back the token it was started with: resuming would loop + // forever without progress. + if (error != null) { + throw rethrow(fail(error)); + } + return false; + } + + noProgressReconnects = + nextNoProgressReconnects(noProgressReconnects, hop.status, hop.startToken, hop.token); + if (noProgressReconnects == 0) { + lastError = null; + } + if (error != null) { + lastError = error; + } + if (noProgressReconnects > InteractionContinuation.MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS) { + if (lastError != null) { + throw rethrow(fail(lastError)); + } + return false; + } + + nextToken = hop.token; + nextTokenOwnerId = hop.tokenOwnerId; + return true; + } + + /** Whether {@code error} is a transport (connection) failure rather than an API error. */ + static boolean isConnectionError(Throwable error) { + for (Throwable t = error; t != null; t = t.getCause()) { + if (t instanceof IOException || t instanceof UncheckedIOException) { + return true; + } + if (t instanceof GaosBaseException) { + return false; + } + if (t.getCause() == t) { + break; + } + } + return false; + } + + /** Rethrows {@code error} as an {@link IOException} or an unchecked exception. */ + private static IOException rethrow(Exception error) throws IOException { + if (error instanceof IOException) { + throw (IOException) error; + } + if (error instanceof RuntimeException) { + throw (RuntimeException) error; + } + throw new IOException(error); + } + + @Nullable + private static String nonEmpty(@Nullable String value) { + return value == null || value.isEmpty() ? null : value; + } +} diff --git a/src/main/java/com/google/genai/gaos/InteractionContinuation.java b/src/main/java/com/google/genai/gaos/InteractionContinuation.java new file mode 100644 index 00000000000..1e9bf49938e --- /dev/null +++ b/src/main/java/com/google/genai/gaos/InteractionContinuation.java @@ -0,0 +1,766 @@ +/* + * 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 + * + * http://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. + */ + +package com.google.genai.gaos; + +import static com.google.genai.gaos.operations.Operations.AsyncRequestOperation; +import static com.google.genai.gaos.operations.Operations.RequestOperation; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genai.gaos.models.errors.ContinuationInterruptedException; +import com.google.genai.gaos.models.interactions.CreateAgentInteraction; +import com.google.genai.gaos.models.interactions.CreateModelInteraction; +import com.google.genai.gaos.models.interactions.Interaction; +import com.google.genai.gaos.models.interactions.InteractionSSEStreamEvent; +import com.google.genai.gaos.models.operations.CreateInteractionRequest; +import com.google.genai.gaos.models.operations.CreateInteractionRequestBody; +import com.google.genai.gaos.models.operations.CreateInteractionResponse; +import com.google.genai.gaos.operations.CreateInteraction; +import com.google.genai.gaos.operations.Operations; +import com.google.genai.gaos.utils.EventStream; +import com.google.genai.gaos.utils.Headers; +import com.google.genai.gaos.utils.Options; +import com.google.genai.gaos.utils.Utils; +import jakarta.annotation.Nullable; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.function.LongSupplier; +import java.util.logging.Logger; + +/** + * Automatic continuation for {@code interactions.create()}. + * + *

Hand-written (not generated). Every {@code create} entry point of {@link Interactions}, {@link + * AsyncInteractions} and their request builders delegates here. It mirrors the Python SDK's {@code + * _interactions_continuation.py}: a unary loop that follows {@code requires_continuation} hops and + * merges them into one {@link Interaction}, a streaming wrapper ({@link ContinuationEventStream}) + * that yields every hop's events in order and resumes from checkpoints, and a whole-call deadline. + * + *

Behavior summary (see {@link Interactions#create(CreateInteractionRequest, Options)}): + * + *

    + *
  • {@code automaticContinuation} defaults to {@code true}. It is a no-op for {@code + * background=true} and for agent interactions, whose loops are managed by the server. + *
  • Continuation hops send {@code continuation_token}; with {@code store} enabled (the default) + * they also send the hop that issued the token as {@code previous_interaction_id} and omit + * {@code input}; with {@code store=false} the original {@code input} is resent. + *
  • A failure after progress throws {@link ContinuationInterruptedException} (cause attached) + * carrying the token to resume from; a first-hop failure throws the underlying error. + *
+ */ +public final class InteractionContinuation { + + private InteractionContinuation() {} + + /** + * Maximum consecutive reconnects after a mid-hop disconnect that made no progress (no new + * checkpoint token since the hop's start token). + */ + public static final int MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS = 3; + + static final String AUTO_CONTINUE_AND_BACKGROUND_ERROR = + "automaticContinuation=false cannot be used with background=true."; + + static final String INCOMPLETE_CONTINUATION_WARNING = + "Interaction stopped before completion (status='requires_continuation') with a" + + " continuation_token because the per-request decode token limit was reached. To resume" + + " decoding automatically, pass automaticContinuation=true (note: each continuation hop" + + " re-prefills and bills input tokens), or use background=true (the most cost-efficient" + + " option for server-managed continuation), or resume manually by passing" + + " continuation_token to interactions.create() (with previous_interaction_id set to this" + + " interaction's id and no input, or with the original input when store=false)."; + + private static final String REQUIRES_CONTINUATION = "requires_continuation"; + private static final String SSE_SENTINEL = "[DONE]"; + + private static final Logger logger = Logger.getLogger(InteractionContinuation.class.getName()); + private static volatile boolean incompleteContinuationWarned = false; + + /** Clock for {@link CallDeadline}, in nanoseconds. Replaceable by tests in this package. */ + static LongSupplier clock = System::nanoTime; + + // Usage list fields are merged per key (e.g. per modality), summing the value. + private static final Map USAGE_LIST_KEYS = new LinkedHashMap<>(); + + static { + USAGE_LIST_KEYS.put("cached_tokens_by_modality", new String[] {"modality", "tokens"}); + USAGE_LIST_KEYS.put("input_tokens_by_modality", new String[] {"modality", "tokens"}); + USAGE_LIST_KEYS.put("output_tokens_by_modality", new String[] {"modality", "tokens"}); + USAGE_LIST_KEYS.put("tool_use_tokens_by_modality", new String[] {"modality", "tokens"}); + USAGE_LIST_KEYS.put("grounding_tool_count", new String[] {"type", "count"}); + } + + // SDK-only convenience fields recomputed from the merged steps by the Interaction constructor. + private static final List OUTPUT_PROPERTIES = + Arrays.asList("output_text", "output_image", "output_audio", "output_video"); + + // --------------------------------------------------------------------------------------------- + // Entry points + // --------------------------------------------------------------------------------------------- + + /** + * Runs a synchronous {@code create} with automatic continuation. + * + * @throws IllegalArgumentException if {@code automaticContinuation=false} is combined with {@code + * background=true} (before any request is sent) + */ + public static CreateInteractionResponse create( + SDKConfiguration sdkConfiguration, + CreateInteractionRequest request, + @Nullable Options options, + Headers headers) { + Prepared prepared = Prepared.of(request); + CallDeadline deadline = new CallDeadline(prepared.timeout, prepared.enableContinuation); + HopSender sender = new HopSender(sdkConfiguration, options, headers); + + CreateInteractionResponse response = sender.send(prepared.request); + if (prepared.stream) { + if (!isEventStream(response)) { + return response; + } + return response.withEvents(wrapStream(sender, prepared, deadline, response.events())); + } + if (prepared.enableContinuation) { + return continueUnary(sender, prepared, deadline, response); + } + if (prepared.callerToken == null) { + response.interaction().ifPresent(InteractionContinuation::warnIfIncompleteUnary); + } + return response; + } + + /** + * Runs an asynchronous {@code create} with automatic continuation. + * + *

The generated async surface exposes the response as a blocking {@link EventStream} (it is + * meant for {@code stream=true}); continuation hops are fetched on the thread iterating it. + * + * @throws IllegalArgumentException if {@code automaticContinuation=false} is combined with {@code + * background=true} (before any request is sent) + */ + public static CompletableFuture> createAsync( + SDKConfiguration sdkConfiguration, + CreateInteractionRequest request, + @Nullable Options options, + Headers headers) { + // Resolve the (required in some environments) retry scheduler before validating the request, + // as the generated code did. + AsyncRequestOperation< + CreateInteractionRequest, + com.google.genai.gaos.models.operations.async.CreateInteractionResponse> + operation = + new CreateInteraction.Async( + sdkConfiguration, options, sdkConfiguration.retryScheduler(), headers); + final Prepared prepared = Prepared.of(request); + final CallDeadline deadline = new CallDeadline(prepared.timeout, prepared.enableContinuation); + final HopSender sender = new HopSender(sdkConfiguration, options, headers); + + return Operations.relayCancel( + Operations.applyBodyReadAsync( + operation.doRequest(prepared.request), operation::handleResponse) + .thenApplyAsync( + response -> { + EventStream events = + new EventStream( + response.rawResponse().body(), + new TypeReference() {}, + Utils.mapper(), + Optional.of(SSE_SENTINEL)); + if (!prepared.stream) { + return events; + } + return wrapStream(sender, prepared, deadline, events); + }, + Operations.streamCompletionExecutor()), + operation); + } + + private static EventStream wrapStream( + final HopSender sender, + final Prepared prepared, + CallDeadline deadline, + EventStream events) { + if (!prepared.enableContinuation) { + if (prepared.callerToken == null) { + return ContinuationEventStream.withIncompleteWarning(events); + } + return events; + } + return new ContinuationEventStream( + events, + (token, tokenOwnerId) -> { + CreateInteractionResponse next = sender.send(prepared.hopRequest(token, tokenOwnerId)); + return isEventStream(next) ? next.events() : null; + }, + prepared.callerToken, + deadline, + prepared.store); + } + + private static boolean isEventStream(CreateInteractionResponse response) { + return !response.interaction().isPresent() + && Utils.contentTypeMatches(response.contentType(), "text/event-stream"); + } + + /** Sends one hop (one HTTP request) through the generated operation. */ + private static final class HopSender { + private final SDKConfiguration sdkConfiguration; + private final Options options; + private final Headers headers; + + HopSender(SDKConfiguration sdkConfiguration, @Nullable Options options, Headers headers) { + this.sdkConfiguration = sdkConfiguration; + this.options = options; + this.headers = headers; + } + + CreateInteractionResponse send(CreateInteractionRequest request) { + RequestOperation operation = + new CreateInteraction.Sync(sdkConfiguration, options, headers); + return operation.handleResponse(operation.doRequest(request)); + } + } + + // --------------------------------------------------------------------------------------------- + // Request preparation + // --------------------------------------------------------------------------------------------- + + /** A validated {@code create} request plus what the continuation loop needs to know about it. */ + static final class Prepared { + /** The initial request. */ + final CreateInteractionRequest request; + + final boolean stream; + final boolean store; + final boolean enableContinuation; + @Nullable final String callerToken; + @Nullable final Duration timeout; + + private Prepared( + CreateInteractionRequest request, + boolean stream, + boolean store, + boolean enableContinuation, + @Nullable String callerToken, + @Nullable Duration timeout) { + this.request = request; + this.stream = stream; + this.store = store; + this.enableContinuation = enableContinuation; + this.callerToken = callerToken; + this.timeout = timeout; + } + + static Prepared of(CreateInteractionRequest request) { + Utils.checkNotNull(request, "request"); + CreateInteractionRequestBody body = + request + .body() + .orElseThrow(() -> new IllegalArgumentException("Request body is required")); + Boolean automaticContinuation = request.automaticContinuation().orElse(null); + boolean background = bodyBool(body, "background"); + boolean stream = bodyBool(body, "stream"); + boolean store = !Boolean.FALSE.equals(bodyBoolObj(body, "store")); + String callerToken = nonEmpty(bodyContinuationToken(body)); + if (background && Boolean.FALSE.equals(automaticContinuation)) { + throw new IllegalArgumentException(AUTO_CONTINUE_AND_BACKGROUND_ERROR); + } + boolean enableContinuation = + !Boolean.FALSE.equals(automaticContinuation) + && !background + && !body.createAgentInteraction().isPresent(); + return new Prepared( + request, stream, store, enableContinuation, callerToken, request.timeout().orElse(null)); + } + + /** + * The request for a continuation hop resuming from {@code token}. + * + *

With {@code store} enabled (the default) the server rebuilds the context from the stored + * interaction that issued {@code token}, so the hop sends that interaction as {@code + * previous_interaction_id} (the server checks the token belongs to it) and omits {@code + * input}. With {@code store=false} nothing is stored, so the original {@code input} is resent + * with the token instead. + * + * @param tokenOwnerId the id of the hop that issued {@code token}, or null when the token did + * not come from an earlier hop (e.g. the caller's own {@code continuation_token}); the + * original request is then resent as-is with {@code token} + */ + CreateInteractionRequest hopRequest(String token, @Nullable String tokenOwnerId) { + CreateInteractionRequestBody body = request.body().get(); + boolean stateful = tokenOwnerId != null && store; + CreateInteractionRequestBody hopBody; + if (body.createModelInteraction().isPresent()) { + CreateModelInteraction copy = copy(body.createModelInteraction().get()); + copy.withContinuationToken(token); + if (stateful) { + copy.withPreviousInteractionId(tokenOwnerId).withInput(null); + } + hopBody = CreateInteractionRequestBody.of(copy); + } else if (body.createAgentInteraction().isPresent()) { + CreateAgentInteraction copy = copy(body.createAgentInteraction().get()); + copy.withContinuationToken(token); + if (stateful) { + copy.withPreviousInteractionId(tokenOwnerId).withInput(null); + } + hopBody = CreateInteractionRequestBody.of(copy); + } else { + hopBody = body; + } + return new CreateInteractionRequest(request.apiVersion().orElse(null), hopBody); + } + } + + /** A boolean field of either body variant ({@code background}, {@code stream}, {@code store}). */ + @Nullable + private static Boolean bodyBoolObj(CreateInteractionRequestBody body, String name) { + if (body.createModelInteraction().isPresent()) { + CreateModelInteraction m = body.createModelInteraction().get(); + switch (name) { + case "background": + return m.background().orElse(null); + case "stream": + return m.stream().orElse(null); + case "store": + return m.store().orElse(null); + default: + return null; + } + } + if (body.createAgentInteraction().isPresent()) { + CreateAgentInteraction a = body.createAgentInteraction().get(); + switch (name) { + case "background": + return a.background().orElse(null); + case "stream": + return a.stream().orElse(null); + case "store": + return a.store().orElse(null); + default: + return null; + } + } + JsonNode json = body.asJson().orElse(null); + JsonNode node = json == null ? null : json.get(name); + return node != null && node.isBoolean() ? node.asBoolean() : null; + } + + private static boolean bodyBool(CreateInteractionRequestBody body, String name) { + return Boolean.TRUE.equals(bodyBoolObj(body, name)); + } + + /** The body's {@code continuation_token}, if any. */ + @Nullable + private static String bodyContinuationToken(CreateInteractionRequestBody body) { + if (body.createModelInteraction().isPresent()) { + return body.createModelInteraction().get().continuationToken().orElse(null); + } + if (body.createAgentInteraction().isPresent()) { + return body.createAgentInteraction().get().continuationToken().orElse(null); + } + JsonNode json = body.asJson().orElse(null); + JsonNode node = json == null ? null : json.get("continuation_token"); + return node != null && node.isTextual() ? node.asText() : null; + } + + /** Deep-copies a model through its JSON form, so the caller's object is never mutated. */ + private static T copy(T value) { + ObjectMapper mapper = Utils.mapper(); + try { + @SuppressWarnings("unchecked") + Class type = (Class) value.getClass(); + return mapper.readValue(mapper.writeValueAsBytes(value), type); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + @Nullable + private static String nonEmpty(@Nullable String value) { + return value == null || value.isEmpty() ? null : value; + } + + // --------------------------------------------------------------------------------------------- + // Whole-call deadline + // --------------------------------------------------------------------------------------------- + + /** The per-call {@code timeout} ran out before automatic continuation finished. */ + @SuppressWarnings("serial") + public static final class ContinuationTimeoutException extends RuntimeException { + ContinuationTimeoutException(String message) { + super(message); + } + } + + /** + * Wall-clock budget of one logical {@code interactions.create()} call. + * + *

A per-call {@code timeout} bounds the whole call rather than each continuation hop: the + * deadline is fixed when the call starts and no further hop starts (and, for streams, no further + * event is read) once it has passed. Per-request HTTP timeouts (e.g. the OkHttp call timeout from + * {@code HttpOptions.timeout}) still apply to each hop separately. + */ + static final class CallDeadline { + @Nullable private final Duration timeout; + private final long deadlineNanos; + private final boolean bounded; + + CallDeadline(@Nullable Duration timeout, boolean enabled) { + this.timeout = timeout; + this.bounded = enabled && timeout != null; + this.deadlineNanos = bounded ? clock.getAsLong() + timeout.toNanos() : 0L; + } + + boolean expired() { + return bounded && clock.getAsLong() - deadlineNanos >= 0; + } + + /** Throws {@link ContinuationTimeoutException} once the deadline has passed. */ + void check() { + if (expired()) { + throw new ContinuationTimeoutException( + "interactions.create() timed out after " + + timeout + + " before automatic continuation finished. Increase the timeout (it bounds the" + + " whole call, across continuation hops) or resume from the continuation_token."); + } + } + } + + // --------------------------------------------------------------------------------------------- + // Warnings + // --------------------------------------------------------------------------------------------- + + static void warnIncompleteContinuation() { + if (!incompleteContinuationWarned) { + incompleteContinuationWarned = true; + logger.warning(INCOMPLETE_CONTINUATION_WARNING); + } + } + + static void resetIncompleteContinuationWarningForTesting() { + incompleteContinuationWarned = false; + } + + /** Emits a one-time warning when unary {@code create()} stops uncontinued. */ + static void warnIfIncompleteUnary(Interaction interaction) { + if (continuationToken(interaction) != null) { + warnIncompleteContinuation(); + } + } + + // --------------------------------------------------------------------------------------------- + // Unary loop + // --------------------------------------------------------------------------------------------- + + /** The token to continue a unary interaction from, or null if it should not be continued. */ + @Nullable + static String continuationToken(Interaction interaction) { + String status = interaction.status().map(s -> s.value()).orElse(null); + if (!REQUIRES_CONTINUATION.equals(status)) { + return null; + } + return nonEmpty(interaction.continuationToken().orElse(null)); + } + + /** The token to continue from after {@code hop}, or null to stop. */ + @Nullable + private static String nextUnaryHopToken(Interaction hop, @Nullable String sentToken) { + String token = continuationToken(hop); + if (token != null && token.equals(sentToken)) { + // The hop handed back the token it was started with: continuing would loop forever + // without progress. + return null; + } + return token; + } + + /** + * Continues a unary interaction across {@code requires_continuation} hops. + * + *

A hop that fails after an earlier hop completed throws {@link + * ContinuationInterruptedException} (with the failure as cause) carrying the hops merged so far + * and the token to resume from, so the output already produced (and billed) is not lost. Running + * out of {@code deadline} before a hop starts is such a failure. + */ + private static CreateInteractionResponse continueUnary( + HopSender sender, + Prepared prepared, + CallDeadline deadline, + CreateInteractionResponse initialResponse) { + if (!initialResponse.interaction().isPresent()) { + return initialResponse; + } + List hops = new ArrayList<>(); + hops.add(initialResponse.interaction().get()); + CreateInteractionResponse lastResponse = initialResponse; + String sentToken = prepared.callerToken; + while (true) { + Interaction last = hops.get(hops.size() - 1); + String token = nextUnaryHopToken(last, sentToken); + if (token == null) { + return lastResponse.withInteraction(mergeUnaryHops(hops)); + } + CreateInteractionResponse response; + try { + deadline.check(); + response = sender.send(prepared.hopRequest(token, nonEmpty(last.id().orElse(null)))); + } catch (RuntimeException e) { + throw interruptedUnary(e, hops, token, prepared.store); + } + if (!response.interaction().isPresent()) { + return response; + } + hops.add(response.interaction().get()); + lastResponse = response; + sentToken = token; + } + } + + /** Builds the error for a continuation hop that failed after {@code hops}. */ + static ContinuationInterruptedException interruptedUnary( + Throwable error, List hops, String token, boolean store) { + Interaction partial = mergeUnaryHops(hops); + String interactionId = store ? nonEmpty(hops.get(hops.size() - 1).id().orElse(null)) : null; + return new ContinuationInterruptedException( + "Automatic continuation was interrupted after " + + hops.size() + + " completed hop(s): " + + error + + ". The partial interaction (status 'requires_continuation') and its" + + " continuation_token can be used to resume.", + error, + partial, + token, + interactionId, + mergeContinuationIds(hops)); + } + + // --------------------------------------------------------------------------------------------- + // Unary merge (on the JSON tree, so fields added to the schema later need no code change) + // --------------------------------------------------------------------------------------------- + + /** + * Merges automatic continuation hops into a single interaction. + * + *

Each hop is its own interaction on the server, and only contains the output generated + * during that hop. The merged result is the final hop (so {@code id}, {@code status}, etc. refer + * to it) with the steps of every hop concatenated, usage summed across hops, and {@code + * continuation_ids} listing every hop's id. Fields describing the turn's request ({@code + * previous_interaction_id}, {@code created}, {@code input}) come from the first hop: later hops + * point back at the hop before them, not at the caller's previous turn. + */ + static Interaction mergeUnaryHops(List hops) { + Interaction last = hops.get(hops.size() - 1); + if (!last.steps().isPresent()) { + return last; + } + ObjectMapper mapper = Utils.mapper(); + List trees = new ArrayList<>(); + for (Interaction hop : hops) { + trees.add((ObjectNode) mapper.valueToTree(hop)); + } + ObjectNode merged = trees.get(trees.size() - 1).deepCopy(); + + ArrayNode steps = mapper.createArrayNode(); + for (int i = 0; i < trees.size(); i++) { + JsonNode hopSteps = trees.get(i).get("steps"); + if (hopSteps == null || !hopSteps.isArray()) { + continue; + } + int start = 0; + if (i > 0) { + // Later hops continue the same turn; a repeated user_input step would duplicate the + // input and cut output_text (which stops at the latest user_input step) short. + while (start < hopSteps.size() + && "user_input".equals(hopSteps.get(start).path("type").asText(null))) { + start++; + } + } + for (int j = start; j < hopSteps.size(); j++) { + steps.add(hopSteps.get(j)); + } + } + merged.set("steps", steps); + + ArrayNode continuationIds = mapper.createArrayNode(); + for (String id : mergeContinuationIds(hops)) { + continuationIds.add(id); + } + merged.set("continuation_ids", continuationIds); + + if (trees.size() > 1) { + ObjectNode first = trees.get(0); + for (String name : Arrays.asList("previous_interaction_id", "created", "input")) { + JsonNode value = first.get(name); + if (value == null || value.isNull()) { + merged.remove(name); + } else { + merged.set(name, value); + } + } + } + + List usages = new ArrayList<>(); + for (ObjectNode tree : trees) { + JsonNode usage = tree.get("usage"); + if (usage != null && usage.isObject()) { + usages.add(usage); + } + } + if (!usages.isEmpty()) { + JsonNode firstUsage = trees.get(0).get("usage"); + merged.set( + "usage", + mergeUsages(usages, firstUsage != null && firstUsage.isObject() ? firstUsage : null)); + } + + // Recomputed from the merged steps by the Interaction constructor. + for (String name : OUTPUT_PROPERTIES) { + merged.remove(name); + } + try { + return mapper.treeToValue(merged, Interaction.class); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + /** Each hop's server {@code continuation_ids} followed by its id, de-duplicated in order. */ + static List mergeContinuationIds(List hops) { + LinkedHashSet ids = new LinkedHashSet<>(); + for (Interaction hop : hops) { + if (hop.continuationIds().isPresent()) { + for (String id : hop.continuationIds().get()) { + if (id != null && !id.isEmpty()) { + ids.add(id); + } + } + } + String id = nonEmpty(hop.id().orElse(null)); + if (id != null) { + ids.add(id); + } + } + return new ArrayList<>(ids); + } + + /** Sums token usage across hops, keeping the first hop's input tokens. */ + static ObjectNode mergeUsages(List usages) { + return mergeUsages(usages, usages.get(0)); + } + + /** Sums token usage across hops, keeping {@code firstHopUsage}'s input tokens. */ + static ObjectNode mergeUsages(List usages, @Nullable JsonNode firstHopUsage) { + ObjectMapper mapper = Utils.mapper(); + ObjectNode merged = usages.get(usages.size() - 1).deepCopy(); + + LinkedHashSet names = new LinkedHashSet<>(); + for (JsonNode usage : usages) { + for (Iterator it = usage.fieldNames(); it.hasNext(); ) { + names.add(it.next()); + } + } + for (String name : names) { + if (name.equals("initial_input_tokens")) { + continue; + } + List present = new ArrayList<>(); + for (JsonNode usage : usages) { + JsonNode value = usage.get(name); + if (value != null && !value.isNull()) { + present.add(value); + } + } + if (present.isEmpty()) { + continue; + } + boolean allInts = true; + boolean allLists = true; + for (JsonNode value : present) { + allInts &= value.isIntegralNumber(); + allLists &= value.isArray(); + } + if (allInts) { + long sum = 0; + for (JsonNode value : present) { + sum += value.asLong(); + } + merged.put(name, sum); + } else if (allLists && USAGE_LIST_KEYS.containsKey(name)) { + String[] keys = USAGE_LIST_KEYS.get(name); + merged.set(name, mergeUsageList(mapper, present, keys[0], keys[1])); + } + } + + JsonNode initialInput = + firstHopUsage == null ? null : firstHopUsage.get("initial_input_tokens"); + if (initialInput == null || initialInput.isNull()) { + initialInput = firstHopUsage == null ? null : firstHopUsage.get("total_input_tokens"); + } + if (initialInput != null && !initialInput.isNull()) { + merged.set("initial_input_tokens", initialInput); + } else { + merged.remove("initial_input_tokens"); + } + return merged; + } + + private static ArrayNode mergeUsageList( + ObjectMapper mapper, List lists, String keyName, String valueName) { + LinkedHashMap merged = new LinkedHashMap<>(); + for (JsonNode items : lists) { + for (JsonNode item : items) { + if (!item.isObject()) { + continue; + } + String key = item.path(keyName).asText(null); + ObjectNode prev = merged.get(key); + if (prev == null) { + merged.put(key, ((ObjectNode) item).deepCopy()); + continue; + } + JsonNode prevValue = prev.get(valueName); + JsonNode value = item.get(valueName); + if (prevValue != null + && prevValue.isIntegralNumber() + && value != null + && value.isIntegralNumber()) { + prev.put(valueName, prevValue.asLong() + value.asLong()); + } + } + } + ArrayNode result = mapper.createArrayNode(); + for (ObjectNode item : merged.values()) { + result.add(item); + } + return result; + } +} diff --git a/src/main/java/com/google/genai/gaos/Interactions.java b/src/main/java/com/google/genai/gaos/Interactions.java index a1280f032ce..708bb61f014 100644 --- a/src/main/java/com/google/genai/gaos/Interactions.java +++ b/src/main/java/com/google/genai/gaos/Interactions.java @@ -35,7 +35,6 @@ import com.google.genai.gaos.models.operations.GetInteractionByIdRequestBuilder; import com.google.genai.gaos.models.operations.GetInteractionByIdResponse; import com.google.genai.gaos.operations.CancelInteractionById; -import com.google.genai.gaos.operations.CreateInteraction; import com.google.genai.gaos.operations.DeleteInteraction; import com.google.genai.gaos.operations.GetInteractionById; import com.google.genai.gaos.utils.Headers; @@ -84,6 +83,21 @@ public CreateInteractionResponse create(@Nonnull CreateInteractionRequestBody bo return create(null, body, null); } + /** + * Creates a new interaction. + * + * @param body The request body. + * @param automaticContinuation Whether to automatically continue decoding across + * continuation hops; see {@link #create(CreateInteractionRequest, Options)}. + * @return The response from the API call + * @throws RuntimeException subclass if the API call fails + */ + public CreateInteractionResponse create( + @Nonnull CreateInteractionRequestBody body, boolean automaticContinuation) { + return create( + new CreateInteractionRequest(null, body, automaticContinuation, null), null); + } + /** * Creates a new interaction. * @@ -96,10 +110,90 @@ public CreateInteractionResponse create(@Nonnull CreateInteractionRequestBody bo public CreateInteractionResponse create( @Nullable String apiVersion, @Nonnull CreateInteractionRequestBody body, @Nullable Options options) { - CreateInteractionRequest request = new CreateInteractionRequest(apiVersion, body); - RequestOperation operation - = new CreateInteraction.Sync(sdkConfiguration, options, _headers); - return operation.handleResponse(operation.doRequest(request)); + return create(new CreateInteractionRequest(apiVersion, body), options); + } + + /** + * Creates a new interaction. + * + * @param request The request object containing all the parameters for the API call. + * @return The response from the API call + * @throws RuntimeException subclass if the API call fails + */ + public CreateInteractionResponse create(@Nonnull CreateInteractionRequest request) { + return create(request, null); + } + + /** + * Creates a new interaction. + * + *

Automatic continuation. {@code request.automaticContinuation()} defaults to + * {@code true}. When enabled, the SDK automatically resumes decoding across multiple + * continuation hops whenever a response or stream stops at the per-request decode token + * limit ({@code status='requires_continuation'} with a {@code continuation_token}) until the + * interaction completes, reaches {@code incomplete} (e.g. {@code max_output_tokens}), + * requires action, or fails. Pass {@code automaticContinuation(false)} to disable; the + * server response is then returned as-is (field for field, continuation fields included) so + * the caller can resume manually, and a one-time warning is logged if it stopped at + * {@code requires_continuation}. + * + *

    + *
  • For {@code stream=true}, events from all continuation hops are yielded sequentially + * on the returned {@link CreateInteractionResponse#events()} stream, verbatim: each hop's + * own {@code interaction.created} and {@code interaction.completed} (with status + * {@code requires_continuation}) are included, so check {@code status} on + * {@code interaction.completed} rather than treating the first one as the end. + *
  • For {@code stream=false}, each continuation hop is a separate interaction on the + * server that only contains that hop's output. The SDK merges the hops into one + * {@code Interaction}: {@code steps} are concatenated in hop order, {@code usage} token + * counts are summed across hops (with {@code usage.initial_input_tokens} set to the + * first hop's input tokens), and all other fields ({@code id}, {@code status}, ...) + * come from the final hop. The final hop's {@code id} can be passed as + * {@code previous_interaction_id} to continue the conversation. + * {@code continuation_ids} lists the ids of every hop in order (the last one is + * {@code id}); note that {@code get(id)} and {@code delete(id)} act on a single hop. + * {@code previous_interaction_id}, {@code created} and {@code input} come from the + * first hop. + *
  • With {@code store=true} (the default), each continuation hop sends the hop that + * issued the {@code continuation_token} as {@code previous_interaction_id} and omits + * {@code input} (the server already has it). With {@code store=false}, the original + * {@code input} is resent with the token. + *
  • Client-side continuation is skipped for {@code background=true} and agent + * interactions ({@code agent} / {@code agent_config}), whose loops are managed by the + * server. {@code background=true} with {@code automaticContinuation(false)} is rejected + * with an {@link IllegalArgumentException} before any request is sent. + *
  • Timeouts: {@code request.timeout()} is a wall-clock budget for the whole + * {@code create()} call. Once it runs out no further hop starts (and no further stream + * event is read) and the call stops with {@code ContinuationInterruptedException} (see + * Errors) so the work so far can be resumed. Per-request HTTP timeouts (e.g. + * {@code HttpOptions.timeout} on the client) and {@code RetryConfig} apply to each hop + * separately: transient errors on a continuation hop retry that hop using the latest + * {@code continuation_token} without restarting from the initial request. + *
  • Billing note: with {@code background=false}, each continuation hop re-prefills the + * conversation and generated prefix, so input tokens are billed on every hop (and + * summed in {@code usage.total_input_tokens}). If an unrecoverable error occurs after + * one or more hops have completed, those completed hops are still billed. For + * server-managed continuation without client-side hop loops, use + * {@code background=true} (with {@code store=true}). + *
  • Errors: a failure on the first request is thrown as usual. A failure after progress + * (a later hop, or a stream after a resumable checkpoint) throws + * {@link com.google.genai.gaos.models.errors.ContinuationInterruptedException} (with + * the failure as cause) carrying {@code continuationToken}, {@code interactionId} and + * {@code continuationIds}, plus {@code partialInteraction} (the completed hops merged, + * for {@code stream=false}), so the output already produced is not lost and the call + * can be resumed manually. Running out of the per-call timeout after progress is + * reported the same way. + *
+ * + * @param request The request object containing all the parameters for the API call. + * @param options additional options + * @return The response from the API call + * @throws RuntimeException subclass if the API call fails + */ + public CreateInteractionResponse create( + @Nonnull CreateInteractionRequest request, + @Nullable Options options) { + return InteractionContinuation.create(sdkConfiguration, request, options, _headers); } /** diff --git a/src/main/java/com/google/genai/gaos/models/errors/ContinuationInterruptedException.java b/src/main/java/com/google/genai/gaos/models/errors/ContinuationInterruptedException.java new file mode 100644 index 00000000000..a21f28494f0 --- /dev/null +++ b/src/main/java/com/google/genai/gaos/models/errors/ContinuationInterruptedException.java @@ -0,0 +1,94 @@ +/* + * 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 + * + * http://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. + */ + +package com.google.genai.gaos.models.errors; + +import com.google.genai.gaos.models.interactions.Interaction; +import jakarta.annotation.Nullable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +/** + * Thrown when Interactions automatic continuation fails after progress. + * + *

{@code interactions.create()} with automatic continuation enabled (the default) may make + * several requests ("hops"), each one resuming the decode from the previous hop's {@code + * continuation_token}. If a hop fails after the server already produced (and billed) output in an + * earlier hop, or after a resumable checkpoint was reached, this exception is thrown instead of + * the underlying error so the work is not lost. The underlying error is {@link #getCause()}. + * + *

To resume, pass {@link #continuationToken()} as {@code continuation_token} to {@code + * interactions.create()}: with {@code store} enabled (the default), also pass {@link + * #interactionId()} as {@code previous_interaction_id} and no {@code input}; with {@code + * store=false}, resend the original {@code input}. + */ +@SuppressWarnings("serial") +public class ContinuationInterruptedException extends RuntimeException { + + private final transient Interaction partialInteraction; + private final String continuationToken; + private final String interactionId; + private final List continuationIds; + + public ContinuationInterruptedException( + String message, + @Nullable Throwable cause, + @Nullable Interaction partialInteraction, + @Nullable String continuationToken, + @Nullable String interactionId, + @Nullable List continuationIds) { + super(message, cause); + this.partialInteraction = partialInteraction; + this.continuationToken = continuationToken; + this.interactionId = interactionId; + this.continuationIds = + continuationIds == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(continuationIds)); + } + + /** + * For unary calls, the interaction merged from the hops that completed (status {@code + * requires_continuation}, with the resumable {@code continuation_token}). Always empty for + * streams, whose events were already yielded. + */ + public Optional partialInteraction() { + return Optional.ofNullable(partialInteraction); + } + + /** + * The latest resumable {@code continuation_token}. For streams this may be a mid-hop checkpoint: + * resuming from it re-generates any output yielded after the checkpoint. + */ + public Optional continuationToken() { + return Optional.ofNullable(continuationToken); + } + + /** + * The id of the interaction that issued {@link #continuationToken()} (the value to send as + * {@code previous_interaction_id}), or empty if unknown or unstored ({@code store=false}). + */ + public Optional interactionId() { + return Optional.ofNullable(interactionId); + } + + /** The ids of every hop created during the call, in order. */ + public List continuationIds() { + return continuationIds; + } +} diff --git a/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequest.java b/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequest.java index be8d84521b7..2d458da9d17 100644 --- a/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequest.java +++ b/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequest.java @@ -26,6 +26,7 @@ import jakarta.annotation.Nullable; import java.lang.Override; import java.lang.String; +import java.time.Duration; import java.util.Optional; @@ -42,13 +43,35 @@ public class CreateInteractionRequest { @SpeakeasyMetadata("request:mediaType=application/json") private CreateInteractionRequestBody body; + /** + * Client-side only (not sent to the API): whether to automatically continue decoding across + * continuation hops. Defaults to true when unset. + */ + private Boolean automaticContinuation; + + /** + * Client-side only (not sent to the API): wall-clock budget for the whole call, across all + * continuation hops. + */ + private Duration timeout; + @JsonCreator public CreateInteractionRequest( @Nullable String apiVersion, @Nonnull CreateInteractionRequestBody body) { + this(apiVersion, body, null, null); + } + + public CreateInteractionRequest( + @Nullable String apiVersion, + @Nonnull CreateInteractionRequestBody body, + @Nullable Boolean automaticContinuation, + @Nullable Duration timeout) { this.apiVersion = apiVersion; this.body = Optional.ofNullable(body) .orElseThrow(() -> new IllegalArgumentException("body cannot be null")); + this.automaticContinuation = automaticContinuation; + this.timeout = timeout; } public CreateInteractionRequest( @@ -70,6 +93,22 @@ public Optional body() { return Optional.ofNullable(this.body); } + /** + * Whether automatic continuation is enabled for this call (defaults to true when empty). + * Client-side only; see {@code Interactions.create}. + */ + public Optional automaticContinuation() { + return Optional.ofNullable(this.automaticContinuation); + } + + /** + * Wall-clock budget for the whole call, across all continuation hops. Client-side only; see + * {@code Interactions.create}. + */ + public Optional timeout() { + return Optional.ofNullable(this.timeout); + } + public static Builder builder() { return new Builder(); } @@ -93,6 +132,26 @@ public CreateInteractionRequest withBody(@Nonnull CreateInteractionRequestBody b } + /** + * Whether automatic continuation is enabled for this call (defaults to true when unset). + * Client-side only; see {@code Interactions.create}. + */ + public CreateInteractionRequest withAutomaticContinuation(@Nullable Boolean automaticContinuation) { + this.automaticContinuation = automaticContinuation; + return this; + } + + + /** + * Wall-clock budget for the whole call, across all continuation hops. Client-side only; see + * {@code Interactions.create}. + */ + public CreateInteractionRequest withTimeout(@Nullable Duration timeout) { + this.timeout = timeout; + return this; + } + + @Override public boolean equals(java.lang.Object o) { if (this == o) { @@ -104,20 +163,25 @@ public boolean equals(java.lang.Object o) { CreateInteractionRequest other = (CreateInteractionRequest) o; return Utils.enhancedDeepEquals(this.apiVersion, other.apiVersion) && - Utils.enhancedDeepEquals(this.body, other.body); + Utils.enhancedDeepEquals(this.body, other.body) && + Utils.enhancedDeepEquals(this.automaticContinuation, other.automaticContinuation) && + Utils.enhancedDeepEquals(this.timeout, other.timeout); } @Override public int hashCode() { return Utils.enhancedHash( - apiVersion, body); + apiVersion, body, automaticContinuation, + timeout); } @Override public String toString() { return Utils.toString(CreateInteractionRequest.class, "apiVersion", apiVersion, - "body", body); + "body", body, + "automaticContinuation", automaticContinuation, + "timeout", timeout); } @SuppressWarnings("UnusedReturnValue") @@ -127,6 +191,10 @@ public final static class Builder { private CreateInteractionRequestBody body; + private Boolean automaticContinuation; + + private Duration timeout; + private Builder() { // force use of static builder() method } @@ -147,9 +215,28 @@ public Builder body(@Nonnull CreateInteractionRequestBody body) { return this; } + /** + * Whether automatic continuation is enabled for this call (defaults to true when unset). + * Client-side only; see {@code Interactions.create}. + */ + public Builder automaticContinuation(@Nullable Boolean automaticContinuation) { + this.automaticContinuation = automaticContinuation; + return this; + } + + /** + * Wall-clock budget for the whole call, across all continuation hops. Client-side only; + * see {@code Interactions.create}. + */ + public Builder timeout(@Nullable Duration timeout) { + this.timeout = timeout; + return this; + } + public CreateInteractionRequest build() { return new CreateInteractionRequest( - apiVersion, body); + apiVersion, body, automaticContinuation, + timeout); } } diff --git a/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequestBuilder.java b/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequestBuilder.java index 6f7bfae7440..c7bba079e26 100644 --- a/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequestBuilder.java +++ b/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionRequestBuilder.java @@ -19,10 +19,8 @@ */ package com.google.genai.gaos.models.operations; -import static com.google.genai.gaos.operations.Operations.RequestOperation; - +import com.google.genai.gaos.InteractionContinuation; import com.google.genai.gaos.SDKConfiguration; -import com.google.genai.gaos.operations.CreateInteraction; import com.google.genai.gaos.utils.Headers; import com.google.genai.gaos.utils.Options; import com.google.genai.gaos.utils.RetryConfig; @@ -30,6 +28,7 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import java.lang.String; +import java.time.Duration; public class CreateInteractionRequestBuilder { private final SDKConfiguration sdkConfiguration; @@ -57,6 +56,26 @@ public CreateInteractionRequestBuilder body(@Nonnull CreateInteractionRequestBod return this; } + /** + * Whether to automatically continue decoding across continuation hops (defaults to true). + * Client-side only, not sent to the API; see {@code Interactions.create}. + */ + public CreateInteractionRequestBuilder automaticContinuation(boolean automaticContinuation) { + this.pojoBuilder.automaticContinuation(automaticContinuation); + this._setterCalled = true; + return this; + } + + /** + * Wall-clock budget for the whole call, across all continuation hops. Client-side only, not + * sent to the API; see {@code Interactions.create}. + */ + public CreateInteractionRequestBuilder timeout(@Nonnull Duration timeout) { + this.pojoBuilder.timeout(Utils.checkNotNull(timeout, "timeout")); + this._setterCalled = true; + return this; + } + public CreateInteractionRequestBuilder retryConfig(RetryConfig retryConfig) { this.optionsBuilder.retryConfig(retryConfig); return this; @@ -83,8 +102,6 @@ public CreateInteractionRequestBuilder header(String name, String value) { */ public CreateInteractionResponse call() { Options options = optionsBuilder.build(); - RequestOperation operation - = new CreateInteraction.Sync(sdkConfiguration, options, _headers); - return operation.handleResponse(operation.doRequest(this._buildRequest())); + return InteractionContinuation.create(sdkConfiguration, this._buildRequest(), options, _headers); } } diff --git a/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionResponse.java b/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionResponse.java index 2dc520b76a0..245cbc52151 100644 --- a/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionResponse.java +++ b/src/main/java/com/google/genai/gaos/models/operations/CreateInteractionResponse.java @@ -110,8 +110,19 @@ public Optional interaction() { } // value overriden by x-speakeasy-sse-sentinel extension via reflection at runtime private Optional _eventSentinel = Optional.empty(); + // Set by automatic continuation to a stream spanning every continuation hop. + private EventStream _events; + /** + * The server-sent events of a {@code stream=true} interaction. + * + *

With automatic continuation enabled (the default), the same stream instance is returned + * on every call and spans every continuation hop; see {@code Interactions.create}. + */ public EventStream events() { + if (_events != null) { + return _events; + } return new EventStream( rawResponse.body(), new TypeReference() {}, @@ -119,6 +130,14 @@ public EventStream events() { _eventSentinel); } + /** + * Replaces the stream returned by {@link #events()} (used by automatic continuation). + */ + public CreateInteractionResponse withEvents(@Nonnull EventStream events) { + this._events = Utils.checkNotNull(events, "events"); + return this; + } + public static Builder builder() { return new Builder(); diff --git a/src/main/java/com/google/genai/gaos/models/operations/async/CreateInteractionRequestBuilder.java b/src/main/java/com/google/genai/gaos/models/operations/async/CreateInteractionRequestBuilder.java index 4929ec11425..e7b685a9405 100644 --- a/src/main/java/com/google/genai/gaos/models/operations/async/CreateInteractionRequestBuilder.java +++ b/src/main/java/com/google/genai/gaos/models/operations/async/CreateInteractionRequestBuilder.java @@ -19,15 +19,11 @@ */ package com.google.genai.gaos.models.operations.async; -import static com.google.genai.gaos.operations.Operations.AsyncRequestOperation; - -import com.fasterxml.jackson.core.type.TypeReference; +import com.google.genai.gaos.InteractionContinuation; import com.google.genai.gaos.SDKConfiguration; import com.google.genai.gaos.models.interactions.InteractionSSEStreamEvent; import com.google.genai.gaos.models.operations.CreateInteractionRequest; import com.google.genai.gaos.models.operations.CreateInteractionRequestBody; -import com.google.genai.gaos.operations.CreateInteraction; -import com.google.genai.gaos.operations.Operations; import com.google.genai.gaos.utils.EventStream; import com.google.genai.gaos.utils.Headers; import com.google.genai.gaos.utils.Options; @@ -36,7 +32,7 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import java.lang.String; -import java.util.Optional; +import java.time.Duration; import java.util.concurrent.CompletableFuture; public class CreateInteractionRequestBuilder { @@ -65,6 +61,26 @@ public CreateInteractionRequestBuilder body(@Nonnull CreateInteractionRequestBod return this; } + /** + * Whether to automatically continue decoding across continuation hops (defaults to true). + * Client-side only, not sent to the API; see {@code Interactions.create}. + */ + public CreateInteractionRequestBuilder automaticContinuation(boolean automaticContinuation) { + this.pojoBuilder.automaticContinuation(automaticContinuation); + this._setterCalled = true; + return this; + } + + /** + * Wall-clock budget for the whole call, across all continuation hops. Client-side only, not + * sent to the API; see {@code Interactions.create}. + */ + public CreateInteractionRequestBuilder timeout(@Nonnull Duration timeout) { + this.pojoBuilder.timeout(Utils.checkNotNull(timeout, "timeout")); + this._setterCalled = true; + return this; + } + public CreateInteractionRequestBuilder retryConfig(RetryConfig retryConfig) { this.optionsBuilder.retryConfig(retryConfig); return this; @@ -98,17 +114,6 @@ public CreateInteractionRequestBuilder header(String name, String value) { */ public CompletableFuture> call() { Options options = optionsBuilder.build(); - AsyncRequestOperation operation - = new CreateInteraction.Async( - sdkConfiguration, options, sdkConfiguration.retryScheduler(), - _headers); - return Operations.relayCancel(Operations.applyBodyReadAsync(operation.doRequest(this._buildRequest()), - operation::handleResponse) - .thenApplyAsync(response -> new EventStream( - response.rawResponse().body(), - new TypeReference() { - }, - Utils.mapper(), - Optional.of("[DONE]")), Operations.streamCompletionExecutor()), operation); + return InteractionContinuation.createAsync(sdkConfiguration, this._buildRequest(), options, _headers); } } diff --git a/src/main/java/com/google/genai/gaos/utils/EventStream.java b/src/main/java/com/google/genai/gaos/utils/EventStream.java index fa6368b6af5..3ee93d862b0 100644 --- a/src/main/java/com/google/genai/gaos/utils/EventStream.java +++ b/src/main/java/com/google/genai/gaos/utils/EventStream.java @@ -82,7 +82,7 @@ * * @param the type that SSE {@code data} fields will be deserialized into */ -public final class EventStream implements Iterable, AutoCloseable { +public class EventStream implements Iterable, AutoCloseable { private static final SpeakeasyLogger logger = SpeakeasyLogger.getLogger(EventStream.class); @@ -94,6 +94,18 @@ public final class EventStream implements Iterable, AutoCloseable { private boolean terminated = false; private boolean closed = false; + /** + * For subclasses that produce events themselves (e.g. by combining several underlying + * streams) and override {@link #next()}, {@link #close()} and {@link #isClosed()}. + */ + protected EventStream() { + this.parser = null; + this.typeReference = null; + this.mapper = null; + this.terminalMessage = Optional.empty(); + this.dataRequired = false; + } + // Internal use only public EventStream(InputStream in, TypeReference typeReference, ObjectMapper mapper, Optional terminalMessage) { this(in, typeReference, mapper, terminalMessage, true); @@ -204,7 +216,9 @@ public Stream stream() { public void close() throws IOException { closed = true; logger.debug("EventStream closed"); - parser.close(); + if (parser != null) { + parser.close(); + } } public boolean isClosed() { diff --git a/src/test/java/com/google/genai/gaos/InteractionContinuationTest.java b/src/test/java/com/google/genai/gaos/InteractionContinuationTest.java new file mode 100644 index 00000000000..a885208c8d4 --- /dev/null +++ b/src/test/java/com/google/genai/gaos/InteractionContinuationTest.java @@ -0,0 +1,904 @@ +/* + * 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. + */ + +package com.google.genai.gaos; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genai.Client; +import com.google.genai.gaos.models.errors.ContinuationInterruptedException; +import com.google.genai.gaos.models.interactions.CreateModelInteraction; +import com.google.genai.gaos.models.interactions.GenerationConfig; +import com.google.genai.gaos.models.interactions.Interaction; +import com.google.genai.gaos.models.interactions.InteractionCompletedEvent; +import com.google.genai.gaos.models.interactions.InteractionCreatedEvent; +import com.google.genai.gaos.models.interactions.InteractionSSEEvent; +import com.google.genai.gaos.models.interactions.InteractionSSEStreamEvent; +import com.google.genai.gaos.models.interactions.InteractionsInput; +import com.google.genai.gaos.models.interactions.StepDelta; +import com.google.genai.gaos.models.operations.CreateInteractionRequest; +import com.google.genai.gaos.models.operations.CreateInteractionRequestBody; +import com.google.genai.gaos.models.operations.CreateInteractionResponse; +import com.google.genai.gaos.utils.EventStream; +import com.google.genai.types.ClientOptions; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.HttpRetryOptions; +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.function.LongSupplier; +import java.util.logging.Handler; +import java.util.logging.LogRecord; +import java.util.logging.Logger; +import okhttp3.Interceptor; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import okio.Buffer; +import okio.BufferedSource; +import okio.Okio; +import okio.Source; +import okio.Timeout; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for Interactions automatic continuation behaviors that the shared conformance spec + * does not cover: the whole-call {@code timeout} (with a fake clock) and the {@code + * CreateInteractionRequest} / builder surfaces. + */ +public final class InteractionContinuationTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final MediaType JSON_MEDIA_TYPE = MediaType.get("application/json"); + private static final MediaType SSE_MEDIA_TYPE = MediaType.get("text/event-stream"); + + /** A manual clock for deadline tests, in nanoseconds. */ + private static final class FakeClock implements LongSupplier { + private long nowNanos = Duration.ofSeconds(1000).toNanos(); + + @Override + public long getAsLong() { + return nowNanos; + } + + void advance(Duration duration) { + nowNanos += duration.toNanos(); + } + } + + private final FakeClock clock = new FakeClock(); + private final LongSupplier realClock = InteractionContinuation.clock; + + @AfterEach + void restoreClock() { + InteractionContinuation.clock = realClock; + } + + private void useFakeClock() { + InteractionContinuation.clock = clock; + } + + // --- Fake server ----------------------------------------------------------------------------- + + /** One scripted HTTP response: a JSON interaction, or SSE events. */ + private static final class Hop { + final String json; + final List sseEvents; + final int advanceClockBeforeEvent; + final Duration advance; + + private Hop(String json, List sseEvents, int advanceClockBeforeEvent, Duration advance) { + this.json = json; + this.sseEvents = sseEvents; + this.advanceClockBeforeEvent = advanceClockBeforeEvent; + this.advance = advance; + } + + static Hop unary(String json, Duration advance) { + return new Hop(json, null, -1, advance); + } + + static Hop sse(List events, int advanceClockBeforeEvent, Duration advance) { + return new Hop(null, events, advanceClockBeforeEvent, advance); + } + } + + private static String unaryJson(String id, String status, String token, String text, int out) { + StringBuilder json = + new StringBuilder("{\"id\":\"") + .append(id) + .append("\",\"status\":\"") + .append(status) + .append("\""); + if (token != null) { + json.append(",\"continuation_token\":\"").append(token).append("\""); + } + json.append(",\"steps\":[{\"type\":\"model_output\",\"content\":[{\"type\":\"text\",\"text\":\"") + .append(text) + .append("\"}]}]"); + json.append(",\"usage\":{\"total_input_tokens\":10,\"total_output_tokens\":") + .append(out) + .append(",\"total_tokens\":") + .append(10 + out) + .append("}}"); + return json.toString(); + } + + private static String created(String id) { + return "event: interaction.created\ndata: {\"event_type\":\"interaction.created\"," + + "\"interaction\":{\"id\":\"" + + id + + "\",\"status\":\"in_progress\"}}\n\n"; + } + + private static String delta(String text) { + return "event: step.delta\ndata: {\"event_type\":\"step.delta\",\"index\":0," + + "\"delta\":{\"type\":\"text\",\"text\":\"" + + text + + "\"}}\n\n"; + } + + private static String completed(String id, String status, String token) { + return "event: interaction.completed\ndata: {\"event_type\":\"interaction.completed\"," + + "\"interaction\":{\"id\":\"" + + id + + "\",\"status\":\"" + + status + + "\"" + + (token == null ? "" : ",\"continuation_token\":\"" + token + "\"") + + "}}\n\n"; + } + + /** Serves scripted hops, recording every request body, and advancing the fake clock. */ + private final class FakeServer implements Interceptor { + private final List hops; + final List requests = new ArrayList<>(); + + FakeServer(Hop... hops) { + this.hops = new ArrayList<>(); + for (Hop hop : hops) { + this.hops.add(hop); + } + } + + @Override + public Response intercept(Chain chain) throws IOException { + Request request = chain.request(); + Buffer buffer = new Buffer(); + request.body().writeTo(buffer); + Hop hop; + synchronized (this) { + requests.add(MAPPER.readTree(buffer.readUtf8())); + assertTrue(requests.size() <= hops.size(), "Unexpected request #" + requests.size()); + hop = hops.get(requests.size() - 1); + } + if (hop.json != null) { + if (hop.advance != null) { + clock.advance(hop.advance); + } + return new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .header("Content-Type", "application/json") + .body(ResponseBody.create(hop.json, JSON_MEDIA_TYPE)) + .build(); + } + // One SSE event per read, so the clock can be advanced right before a given event. + List chunks = new ArrayList<>(); + for (String event : hop.sseEvents) { + chunks.add(new Buffer().writeUtf8(event)); + } + Source source = + new Source() { + int index = 0; + + @Override + public long read(Buffer sink, long byteCount) throws IOException { + while (index < chunks.size() && chunks.get(index).size() == 0) { + index++; + } + if (index >= chunks.size()) { + return -1L; + } + if (index == hop.advanceClockBeforeEvent && hop.advance != null) { + clock.advance(hop.advance); + } + return chunks.get(index).read(sink, byteCount); + } + + @Override + public Timeout timeout() { + return Timeout.NONE; + } + + @Override + public void close() {} + }; + BufferedSource buffered = Okio.buffer(source); + ResponseBody body = + new ResponseBody() { + @Override + public MediaType contentType() { + return SSE_MEDIA_TYPE; + } + + @Override + public long contentLength() { + return -1L; + } + + @Override + public BufferedSource source() { + return buffered; + } + }; + return new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .header("Content-Type", "text/event-stream") + .body(body) + .build(); + } + } + + private static Client makeClient(Interceptor fake) { + return Client.builder() + .apiKey("test-key") + .vertexAI(false) + .httpOptions( + HttpOptions.builder() + .retryOptions(HttpRetryOptions.builder().attempts(1).build()) + .build()) + .clientOptions( + ClientOptions.builder() + .customHttpClient(new OkHttpClient.Builder().addInterceptor(fake).build()) + .build()) + .build(); + } + + private static CreateModelInteraction.Builder modelBody() { + return CreateModelInteraction.builder().model("gemini-2.5-pro").input(InteractionsInput.of("Hello")); + } + + private static String eventType(InteractionSSEStreamEvent event) { + InteractionSSEEvent data = event.data().get(); + if (data instanceof InteractionCreatedEvent) { + return "interaction.created"; + } + if (data instanceof StepDelta) { + return "step.delta"; + } + if (data instanceof InteractionCompletedEvent) { + return "interaction.completed"; + } + return data.eventType(); + } + + // --- CallDeadline ---------------------------------------------------------------------------- + + @Test + public void callDeadlineBoundsTheWholeCallOnlyWhenEnabled() { + useFakeClock(); + InteractionContinuation.CallDeadline deadline = + new InteractionContinuation.CallDeadline(Duration.ofSeconds(10), true); + clock.advance(Duration.ofSeconds(4)); + assertFalse(deadline.expired()); + deadline.check(); + clock.advance(Duration.ofSeconds(7)); + assertTrue(deadline.expired()); + assertThrows(InteractionContinuation.ContinuationTimeoutException.class, deadline::check); + + // No timeout, or continuation disabled (opt-out / background / agent calls): never expires. + InteractionContinuation.CallDeadline unbounded = + new InteractionContinuation.CallDeadline(null, true); + InteractionContinuation.CallDeadline disabled = + new InteractionContinuation.CallDeadline(Duration.ofSeconds(60), false); + clock.advance(Duration.ofDays(365)); + assertFalse(unbounded.expired()); + unbounded.check(); + assertFalse(disabled.expired()); + disabled.check(); + } + + // --- Unary timeout --------------------------------------------------------------------------- + + @Test + public void unaryTimeoutWithinBudgetCompletes() { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.unary(unaryJson("i1", "requires_continuation", "t1", "A ", 40), Duration.ofSeconds(10)), + Hop.unary(unaryJson("i2", "requires_continuation", "t2", "B ", 30), Duration.ofSeconds(10)), + Hop.unary(unaryJson("i3", "completed", null, "C", 20), Duration.ofSeconds(10))); + Client client = makeClient(fake); + try { + Interaction result = + client + .interactions + .create() + .body(CreateInteractionRequestBody.of(modelBody().build())) + .timeout(Duration.ofSeconds(60)) + .call() + .interaction() + .get(); + assertEquals("completed", result.status().get().value()); + assertEquals(Arrays.asList("i1", "i2", "i3"), result.continuationIds().get()); + assertEquals("A B C", result.outputText().get()); + assertEquals(3, fake.requests.size()); + } finally { + client.close(); + } + } + + @Test + public void unaryTimeoutExpiryThrowsInterruptedWithPartialResult() { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.unary(unaryJson("i1", "requires_continuation", "t1", "A ", 40), Duration.ofSeconds(10)), + Hop.unary(unaryJson("i2", "requires_continuation", "t2", "B ", 30), Duration.ofSeconds(10)), + Hop.unary(unaryJson("i3", "completed", null, "C", 20), Duration.ofSeconds(10))); + Client client = makeClient(fake); + try { + CreateInteractionRequest request = + CreateInteractionRequest.builder() + .body(CreateInteractionRequestBody.of(modelBody().build())) + .timeout(Duration.ofSeconds(15)) + .build(); + ContinuationInterruptedException error = + assertThrows( + ContinuationInterruptedException.class, () -> client.interactions.create(request)); + // Hop 1 (10s) left 5s for hop 2, which used it up; hop 3 never started. + assertEquals(2, fake.requests.size()); + assertTrue( + error.getCause() instanceof InteractionContinuation.ContinuationTimeoutException, + String.valueOf(error.getCause())); + assertEquals("t2", error.continuationToken().get()); + assertEquals("i2", error.interactionId().get()); + assertEquals(Arrays.asList("i1", "i2"), error.continuationIds()); + Interaction partial = error.partialInteraction().get(); + assertEquals("requires_continuation", partial.status().get().value()); + assertEquals("A B ", partial.outputText().get()); + assertEquals(70, partial.usage().get().totalOutputTokens().get()); + assertEquals(10, partial.usage().get().initialInputTokens().get()); + } finally { + client.close(); + } + } + + @Test + public void unaryStoreFalseLeavesInteractionIdEmptyOnInterrupted() { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.unary( + unaryJson("i1", "requires_continuation", "t1", "A ", 40), Duration.ofSeconds(10)), + Hop.unary(unaryJson("i2", "completed", null, "B", 20), Duration.ofSeconds(10))); + Client client = makeClient(fake); + try { + CreateInteractionRequest request = + CreateInteractionRequest.builder() + .body(CreateInteractionRequestBody.of(modelBody().store(false).build())) + .timeout(Duration.ofSeconds(5)) + .build(); + ContinuationInterruptedException error = + assertThrows( + ContinuationInterruptedException.class, () -> client.interactions.create(request)); + assertEquals("t1", error.continuationToken().get()); + assertFalse(error.interactionId().isPresent()); + assertEquals(Arrays.asList("i1"), error.continuationIds()); + } finally { + client.close(); + } + } + + @Test + public void timeoutIsIgnoredWhenContinuationIsDisabled() { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.unary( + unaryJson("i1", "requires_continuation", "t1", "A ", 40), Duration.ofDays(1))); + Client client = makeClient(fake); + try { + CreateInteractionResponse response = + client + .interactions + .create() + .body(CreateInteractionRequestBody.of(modelBody().build())) + .automaticContinuation(false) + .timeout(Duration.ofSeconds(1)) + .call(); + assertEquals("requires_continuation", response.interaction().get().status().get().value()); + assertEquals(1, fake.requests.size()); + } finally { + client.close(); + } + } + + // --- Stream timeout -------------------------------------------------------------------------- + + @Test + public void streamTimeoutExpiryBetweenHopsThrowsInterrupted() { + useFakeClock(); + FakeServer fake = + new FakeServer( + // Hop 1 outlives the 15s budget before its terminal event. + Hop.sse( + Arrays.asList(created("i1"), delta("Part 1 "), completed("i1", "requires_continuation", "t1")), + /* advanceClockBeforeEvent= */ 2, + Duration.ofSeconds(20))); + Client client = makeClient(fake); + try { + EventStream stream = + client + .interactions + .create() + .body(CreateInteractionRequestBody.of(modelBody().stream(true).build())) + .timeout(Duration.ofSeconds(15)) + .call() + .events(); + List received = new ArrayList<>(); + ContinuationInterruptedException error = + assertThrows( + ContinuationInterruptedException.class, + () -> { + for (InteractionSSEStreamEvent event : stream) { + received.add(eventType(event)); + } + }); + // Hop 1's events were all yielded; hop 2 never started. + assertEquals( + Arrays.asList("interaction.created", "step.delta", "interaction.completed"), received); + assertEquals(1, fake.requests.size()); + assertTrue( + error.getCause() instanceof InteractionContinuation.ContinuationTimeoutException, + String.valueOf(error.getCause())); + assertEquals("t1", error.continuationToken().get()); + assertEquals("i1", error.interactionId().get()); + assertFalse(error.partialInteraction().isPresent()); + } finally { + client.close(); + } + } + + @Test + public void streamStoreFalseLeavesInteractionIdEmptyOnInterrupted() { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.sse( + Arrays.asList( + created("i1"), delta("Part 1 "), completed("i1", "requires_continuation", "t1")), + /* advanceClockBeforeEvent= */ 2, + Duration.ofSeconds(20))); + Client client = makeClient(fake); + try { + EventStream stream = + client + .interactions + .create() + .body(CreateInteractionRequestBody.of(modelBody().stream(true).store(false).build())) + .timeout(Duration.ofSeconds(15)) + .call() + .events(); + ContinuationInterruptedException error = + assertThrows( + ContinuationInterruptedException.class, + () -> { + for (InteractionSSEStreamEvent ignored : stream) {} + }); + assertTrue( + error.getCause() instanceof InteractionContinuation.ContinuationTimeoutException, + String.valueOf(error.getCause())); + assertEquals("t1", error.continuationToken().get()); + assertFalse(error.interactionId().isPresent()); + assertEquals(Arrays.asList("i1"), error.continuationIds()); + assertFalse(error.partialInteraction().isPresent()); + } finally { + client.close(); + } + } + + @Test + public void streamTimeoutExpiryMidHopThrowsInterruptedWithPreviousToken() { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.sse( + Arrays.asList(created("i1"), delta("Part 1 "), completed("i1", "requires_continuation", "t1")), + -1, + null), + // The deadline passes while waiting for the third event of hop 2. + Hop.sse( + Arrays.asList( + created("i2"), + delta("Part 2a "), + delta("Part 2b "), + completed("i2", "completed", null)), + /* advanceClockBeforeEvent= */ 2, + Duration.ofSeconds(100))); + Client client = makeClient(fake); + try { + EventStream stream = + client + .interactions + .create() + .body(CreateInteractionRequestBody.of(modelBody().stream(true).build())) + .timeout(Duration.ofSeconds(60)) + .call() + .events(); + List received = new ArrayList<>(); + ContinuationInterruptedException error = + assertThrows( + ContinuationInterruptedException.class, + () -> { + for (InteractionSSEStreamEvent event : stream) { + received.add(eventType(event)); + } + }); + // The event that arrived as the deadline passed is still delivered; the stream stops + // before waiting for the next one. + assertEquals( + Arrays.asList( + "interaction.created", + "step.delta", + "interaction.completed", + "interaction.created", + "step.delta", + "step.delta"), + received); + assertEquals(2, fake.requests.size()); + assertTrue( + error.getCause() instanceof InteractionContinuation.ContinuationTimeoutException, + String.valueOf(error.getCause())); + // Hop 2 produced output after hop 1's token, so resuming from t1 re-generates "Part 2a ". + assertEquals("t1", error.continuationToken().get()); + assertEquals("i1", error.interactionId().get()); + assertEquals(Arrays.asList("i1", "i2"), error.continuationIds()); + } finally { + client.close(); + } + } + + // --- Request surface ------------------------------------------------------------------------- + + @Test + public void requestObjectOptOutIsAPassthroughAndDoesNotMutateTheBody() { + FakeServer fake = + new FakeServer(Hop.unary(unaryJson("i1", "requires_continuation", "t1", "A", 5), null)); + Client client = makeClient(fake); + try { + CreateModelInteraction body = modelBody().build(); + CreateInteractionRequest request = + new CreateInteractionRequest(null, CreateInteractionRequestBody.of(body), false, null); + Interaction result = client.interactions.create(request).interaction().get(); + assertEquals(1, fake.requests.size()); + assertFalse(fake.requests.get(0).has("generation_config"), fake.requests.get(0).toString()); + assertEquals("t1", result.continuationToken().get()); + assertFalse(result.continuationIds().isPresent()); + assertFalse(body.generationConfig().isPresent()); + } finally { + client.close(); + } + } + + @Test + public void defaultContinuationLeavesMaxOutputTokensUnset() { + FakeServer fake = + new FakeServer( + Hop.unary(unaryJson("i1", "requires_continuation", "t1", "A ", 5), null), + Hop.unary(unaryJson("i2", "completed", null, "B", 5), null)); + Client client = makeClient(fake); + try { + Interaction result = + client + .interactions + .create(CreateInteractionRequestBody.of(modelBody().build())) + .interaction() + .get(); + assertEquals("completed", result.status().get().value()); + assertEquals("A B", result.outputText().get()); + assertEquals(Arrays.asList("i1", "i2"), result.continuationIds().get()); + assertEquals(2, fake.requests.size()); + assertFalse(fake.requests.get(0).has("generation_config"), fake.requests.get(0).toString()); + assertFalse(fake.requests.get(1).has("generation_config"), fake.requests.get(1).toString()); + } finally { + client.close(); + } + } + + @Test + public void generationConfigWithoutMaxOutputTokensIsPreservedWithoutMutatingCallersObjects() { + FakeServer fake = new FakeServer(Hop.unary(unaryJson("i1", "completed", null, "A", 5), null)); + Client client = makeClient(fake); + try { + GenerationConfig generationConfig = GenerationConfig.builder().seed(7).build(); + CreateModelInteraction body = modelBody().generationConfig(generationConfig).build(); + Interaction result = + client.interactions.create(CreateInteractionRequestBody.of(body)).interaction().get(); + assertEquals("completed", result.status().get().value()); + JsonNode sent = fake.requests.get(0).get("generation_config"); + assertEquals(7, sent.get("seed").asInt()); + assertFalse(sent.has("max_output_tokens"), sent.toString()); + assertFalse(generationConfig.maxOutputTokens().isPresent()); + assertSame(generationConfig, body.generationConfig().get()); + } finally { + client.close(); + } + } + + @Test + public void bodyOverloadWithAutomaticContinuationFlag() { + FakeServer fake = + new FakeServer(Hop.unary(unaryJson("i1", "requires_continuation", "t1", "A", 5), null)); + Client client = makeClient(fake); + try { + Interaction result = + client + .interactions + .create(CreateInteractionRequestBody.of(modelBody().build()), false) + .interaction() + .get(); + assertEquals(1, fake.requests.size()); + assertEquals("requires_continuation", result.status().get().value()); + } finally { + client.close(); + } + } + + @Test + public void streamTimeoutExpiryBeforeProgressThrowsRawTimeout() { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.sse( + Arrays.asList( + created("i1"), + delta("Part 1 "), + completed("i1", "requires_continuation", "t1")), + /* advanceClockBeforeEvent= */ 1, + Duration.ofSeconds(20))); + Client client = makeClient(fake); + try { + EventStream stream = + client + .interactions + .create() + .body(CreateInteractionRequestBody.of(modelBody().stream(true).build())) + .timeout(Duration.ofSeconds(15)) + .call() + .events(); + List received = new ArrayList<>(); + assertThrows( + InteractionContinuation.ContinuationTimeoutException.class, + () -> { + for (InteractionSSEStreamEvent event : stream) { + received.add(eventType(event)); + } + }); + // No checkpoint token had been issued yet, so the raw timeout error is thrown directly. + assertEquals(Arrays.asList("interaction.created", "step.delta"), received); + assertEquals(1, fake.requests.size()); + } finally { + client.close(); + } + } + + @Test + public void asyncStreamTimeoutExpiryAfterProgressThrowsInterrupted() throws Exception { + useFakeClock(); + FakeServer fake = + new FakeServer( + Hop.sse( + Arrays.asList( + created("i1"), + delta("Part 1 "), + completed("i1", "requires_continuation", "t1")), + /* advanceClockBeforeEvent= */ 2, + Duration.ofSeconds(20))); + ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); + Client client = + Client.builder() + .apiKey("test-key") + .vertexAI(false) + .asyncRetryScheduler(scheduler) + .httpOptions( + HttpOptions.builder() + .retryOptions(HttpRetryOptions.builder().attempts(1).build()) + .build()) + .clientOptions( + ClientOptions.builder() + .customHttpClient(new OkHttpClient.Builder().addInterceptor(fake).build()) + .build()) + .build(); + try { + EventStream stream = + client + .async + .interactions + .create() + .body(CreateInteractionRequestBody.of(modelBody().stream(true).build())) + .timeout(Duration.ofSeconds(15)) + .call() + .get(); + ContinuationInterruptedException error = + assertThrows( + ContinuationInterruptedException.class, + () -> { + for (InteractionSSEStreamEvent ignored : stream) {} + }); + assertTrue(error.getCause() instanceof InteractionContinuation.ContinuationTimeoutException); + assertEquals("t1", error.continuationToken().get()); + assertEquals("i1", error.interactionId().get()); + } finally { + client.close(); + scheduler.shutdownNow(); + } + } + + @Test + public void optOutEmitsOneTimeWarningOnlyWithoutCallerToken() throws Exception { + Logger logger = Logger.getLogger(InteractionContinuation.class.getName()); + final List warnings = new ArrayList<>(); + Handler handler = + new Handler() { + @Override + public void publish(LogRecord record) { + warnings.add(record.getMessage()); + } + + @Override + public void flush() {} + + @Override + public void close() {} + }; + logger.addHandler(handler); + InteractionContinuation.resetIncompleteContinuationWarningForTesting(); + FakeServer fake = + new FakeServer( + Hop.unary(unaryJson("i1", "requires_continuation", "t1", "A", 5), null), + Hop.unary(unaryJson("i2", "requires_continuation", "t2", "B", 5), null), + Hop.unary(unaryJson("i3", "requires_continuation", "t3", "C", 5), null), + Hop.sse( + Arrays.asList( + created("i4"), delta("D"), completed("i4", "requires_continuation", "t4")), + -1, + null)); + Client client = makeClient(fake); + try { + // Caller already passing continuation_token -> no warning. + CreateInteractionResponse withToken = + client.interactions.create( + CreateInteractionRequestBody.of(modelBody().continuationToken("caller-tok").build()), + false); + assertTrue(withToken.interaction().isPresent()); + assertEquals(0, warnings.size()); + + // First opt-out unary call stopping at requires_continuation -> warns once. + CreateInteractionResponse firstOptOut = + client.interactions.create(CreateInteractionRequestBody.of(modelBody().build()), false); + assertTrue(firstOptOut.interaction().isPresent()); + assertEquals(1, warnings.size()); + assertEquals(InteractionContinuation.INCOMPLETE_CONTINUATION_WARNING, warnings.get(0)); + + // Second opt-out call -> suppressed (process-wide one-time warning). + CreateInteractionResponse secondOptOut = + client.interactions.create(CreateInteractionRequestBody.of(modelBody().build()), false); + assertTrue(secondOptOut.interaction().isPresent()); + assertEquals(1, warnings.size()); + + // Reset and verify streaming opt-out also warns once. + InteractionContinuation.resetIncompleteContinuationWarningForTesting(); + try (EventStream stream = + client + .interactions + .create(CreateInteractionRequestBody.of(modelBody().stream(true).build()), false) + .events()) { + for (InteractionSSEStreamEvent ignored : stream) {} + } + assertEquals(2, warnings.size()); + } finally { + logger.removeHandler(handler); + client.close(); + } + } + + // --- mergeUsages ----------------------------------------------------------------------------- + + @Test + public void mergeUsagesRetainsHop1InitialInputTokensAndSumsTotals() throws Exception { + ObjectNode merged = + InteractionContinuation.mergeUsages( + Arrays.asList( + MAPPER.readTree( + "{\"total_input_tokens\":10,\"total_output_tokens\":20," + + "\"initial_input_tokens\":10}"), + MAPPER.readTree( + "{\"total_input_tokens\":5,\"total_output_tokens\":30," + + "\"initial_input_tokens\":5}"))); + assertEquals(15, merged.get("total_input_tokens").asInt()); + assertEquals(50, merged.get("total_output_tokens").asInt()); + assertEquals(10, merged.get("initial_input_tokens").asInt()); + } + + @Test + public void mergeUsagesFallsBackToHop1TotalInputTokens() throws Exception { + ObjectNode merged = + InteractionContinuation.mergeUsages( + Arrays.asList( + MAPPER.readTree("{\"total_input_tokens\":10,\"total_output_tokens\":20}"), + MAPPER.readTree( + "{\"total_input_tokens\":5,\"total_output_tokens\":30," + + "\"initial_input_tokens\":5}"))); + assertEquals(15, merged.get("total_input_tokens").asInt()); + assertEquals(50, merged.get("total_output_tokens").asInt()); + assertEquals(10, merged.get("initial_input_tokens").asInt()); + } + + @Test + public void mergeUsagesDoesNotLeakHopNInitialInputTokensWhenHop1LacksInputTokens() + throws Exception { + ObjectNode merged = + InteractionContinuation.mergeUsages( + Arrays.asList( + MAPPER.readTree("{\"total_output_tokens\":20}"), + MAPPER.readTree("{\"total_output_tokens\":30,\"initial_input_tokens\":5}"))); + assertEquals(50, merged.get("total_output_tokens").asInt()); + assertFalse(merged.has("initial_input_tokens"), merged.toString()); + + // Also when Hop 1 omits usage altogether. + Interaction hop1 = + MAPPER.readValue( + "{\"id\":\"i1\",\"status\":\"requires_continuation\",\"continuation_token\":\"t1\"," + + "\"steps\":[{\"type\":\"model_output\"," + + "\"content\":[{\"type\":\"text\",\"text\":\"A\"}]}]}", + Interaction.class); + Interaction hop2 = + MAPPER.readValue( + "{\"id\":\"i2\",\"status\":\"completed\"," + + "\"steps\":[{\"type\":\"model_output\"," + + "\"content\":[{\"type\":\"text\",\"text\":\"B\"}]}]," + + "\"usage\":{\"total_input_tokens\":5,\"total_output_tokens\":30," + + "\"initial_input_tokens\":5}}", + Interaction.class); + Interaction mergedInteraction = + InteractionContinuation.mergeUnaryHops(Arrays.asList(hop1, hop2)); + assertEquals(5, mergedInteraction.usage().get().totalInputTokens().get()); + assertEquals(30, mergedInteraction.usage().get().totalOutputTokens().get()); + assertFalse(mergedInteraction.usage().get().initialInputTokens().isPresent()); + } +}