Add durable conversation resumability to the Antigravity Interactions harness

internal/harness/antigravityinteractions.go

@@ -30,29 +30,35 @@ package harness
 //     no executor is configured, no third-party tools are advertised.
 //
 // Neither kind of tool call is surfaced to the caller: Run drives the whole
-// interaction to completion (initial turn -> resume -> resume -> ... -> final
-// answer) and only forwards the agent's text output via Handler.OnMessage.
+// interaction loop to completion (initial turn -> continuation turn -> ... ->
+// final answer) and only forwards the agent's text output via Handler.OnMessage.
+// Each turn after the first is chained to the previous one via the Interactions
+// API's previous_interaction_id; this is an implementation detail of the loop,
+// not an AX resume.
 //
-// Queue carries human input only -- the initial prompt and, in the future,
-// "steering" messages injected mid-run. It never carries tool results (the
-// harness produces those itself). Queued input is drained at each interaction
-// gap (every resume point), which is the only place the harness can influence an
-// otherwise atomic interaction.
+// Queue carries human input only -- the initial prompt and "steering" messages
+// injected mid-run. It never carries tool results (the harness produces those
+// itself). Queued input is drained at each interaction gap, which is the only
+// place the harness can influence an otherwise atomic interaction.
 
 import (
 	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
+	"math/rand"
 	"net/http"
 	"os"
 	"sort"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/google/ax/internal/storage"
 	"github.com/google/ax/proto"
 	"github.com/google/uuid"
 	"golang.org/x/oauth2"
@@ -91,7 +97,7 @@ type AntigravityInteractionsConfig struct {
 	Agent string
 	// SystemInstruction, if set, is sent as the interaction's system_instruction
 	// (a free-form system prompt prepended to the agent's own instructions). It
-	// is sent on every turn so it persists across resumes.
+	// is sent on every turn of the interaction loop so it persists across them.
 	SystemInstruction string
 	// MaxTurns caps the number of interaction turns the harness will drive within
 	// a single Run before giving up. Defaults to 100.
@@ -115,6 +121,15 @@ type AntigravityInteractionsConfig struct {
 	// HTTPClient overrides the HTTP client. If nil, a default client with a long
 	// timeout is used.
 	HTTPClient *http.Client
+
+	// StateStore durably persists each conversation's resume cursor (the last
+	// interaction id) so a conversation can resume after a process restart or on
+	// a different replica. If nil, that cursor is kept in memory only (on the
+	// Execution) and is lost when the Execution is discarded or the process
+	// exits -- a later Start for the same conversation id then begins a new
+	// interaction chain (empty previous_interaction_id) instead of continuing the
+	// existing one.
+	StateStore storage.Store
 }
 
 func (c *AntigravityInteractionsConfig) withDefaults() {
@@ -162,14 +177,55 @@ func NewAntigravityInteractionsHarness(cfg AntigravityInteractionsConfig) *Antig
 	return &AntigravityInteractionsHarness{cfg: cfg, httpClient: hc}
 }
 
-// Start implements Harness.Start.
+// resumeCursor is the small per-conversation state persisted to StateStore so a
+// conversation can resume across restarts/replicas. It records the tail of the
+// server-side interaction chain (PrevInteractionID). The cursor is only ever
+// persisted after a turn completes successfully, so a persisted cursor always
+// has a non-empty PrevInteractionID; "the conversation has started" is therefore
+// derived as PrevInteractionID != "" rather than stored separately.
+type resumeCursor struct {
+	PrevInteractionID string `json:"prev_interaction_id"`
+}
+
+// stateStoreKey is the StateStore key for a conversation's resume cursor.
+func stateStoreKey(conversationID string) string {
+	return "antigravity-interactions/cursor/" + conversationID
+}
+
+// Start implements Harness.Start. If a StateStore is configured, it loads any
+// previously persisted resume cursor for conversationID so the returned
+// Execution resumes the existing interaction chain instead of starting a new
+// one.
 func (h *AntigravityInteractionsHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (Execution, error) {
-	return &antigravityInteractionsExecution{
+	e := &antigravityInteractionsExecution{
 		harness:        h,
 		conversationID: conversationID,
 		id:             uuid.NewString(),
 		harnessConfig:  harnessConfig,
-	}, nil
+	}
+	if h.cfg.StateStore == nil {
+		return e, nil
+	}
+
+	value, err := h.cfg.StateStore.Get(ctx, stateStoreKey(conversationID))
+	switch {
+	case errors.Is(err, storage.ErrNotFound):
+		// No prior state: a brand-new conversation.
+	case err != nil:
+		// A real storage failure must not be silently treated as "new", which
+		// would lose an existing conversation's history.
+		return nil, fmt.Errorf("loading resume cursor for %q: %w", conversationID, err)
+	default:
+		var cur resumeCursor
+		if err := json.Unmarshal(value, &cur); err != nil {
+			return nil, fmt.Errorf("decoding resume cursor for %q: %w", conversationID, err)
+		}
+		// A persisted cursor is only written after a successful turn, so a
+		// non-empty interaction id means the conversation has already started;
+		// the first-turn check in Run derives that from prevInteractionID.
+		e.prevInteractionID = cur.PrevInteractionID
+	}
+	return e, nil
 }
 
 // antigravityInteractionsExecution implements Execution. It is long-lived
@@ -186,10 +242,10 @@ type antigravityInteractionsExecution struct {
 	queued []*proto.Message
 	closed bool
 
-	// started is false until the initial turn has been sent.
-	started bool
-	// prevInteractionID chains resume turns (the interaction chain this Execution
-	// owns).
+	// prevInteractionID chains the turns of the interaction loop (the interaction
+	// chain this Execution owns), and is the value persisted for cross-Execution
+	// resume. It is empty until the first turn completes successfully; an empty
+	// value therefore means "no turn has succeeded yet" (the first turn).
 	prevInteractionID string
 }
 
@@ -229,13 +285,36 @@ func (e *antigravityInteractionsExecution) drainQueue() []any {
 	return messagesToInputSteps(msgs)
 }
 
-func (e *antigravityInteractionsExecution) setPrevID(id string) {
+// setPrevID records the latest interaction id (in memory) and, if a StateStore
+// is configured, durably persists the resume cursor so the conversation can
+// resume after a restart. A persistence failure is returned to the caller so
+// the run can decide how to handle it rather than silently losing durability.
+func (e *antigravityInteractionsExecution) setPrevID(ctx context.Context, id string) error {
 	if id == "" {
-		return
+		return nil
 	}
 	e.mu.Lock()
 	e.prevInteractionID = id
 	e.mu.Unlock()
+
+	if e.harness.cfg.StateStore == nil {
+		return nil
+	}
+	return e.persistCursor(ctx, resumeCursor{PrevInteractionID: id})
+}
+
+// persistCursor writes the resume cursor to the StateStore. The store is
+// last-write-wins; correctness relies on the controller guaranteeing a single
+// Execution (writer) per conversation (see the Harness interface).
+func (e *antigravityInteractionsExecution) persistCursor(ctx context.Context, cur resumeCursor) error {
+	blob, err := json.Marshal(cur)
+	if err != nil {
+		return fmt.Errorf("encoding resume cursor: %w", err)
+	}
+	if err := e.harness.cfg.StateStore.Put(ctx, stateStoreKey(e.conversationID), blob); err != nil {
+		return fmt.Errorf("persisting resume cursor: %w", err)
+	}
+	return nil
 }
 
 // Run implements Execution.Run. It drives the interaction to completion,
@@ -252,7 +331,6 @@ func (e *antigravityInteractionsExecution) Run(ctx context.Context, handler Hand
 		e.mu.Unlock()
 		return fmt.Errorf("execution session already closed")
 	}
-	started := e.started
 	prevID := e.prevInteractionID
 	e.mu.Unlock()
 
@@ -262,16 +340,15 @@ func (e *antigravityInteractionsExecution) Run(ctx context.Context, handler Hand
 	}
 
 	// Initial input for this Run: drain whatever is queued (the prompt on the
-	// first Run, or steering input on later Runs).
+	// first Run, or steering input on later Runs). An empty prevID means no turn
+	// has completed yet, so this is the conversation's first turn. (A first turn
+	// that failed leaves prevID empty, so a retried Run is correctly treated as
+	// the first turn again.)
 	input := e.drainQueue()
-	if !started {
-		if len(input) == 0 {
+	if len(input) == 0 {
+		if prevID == "" {
 			return fmt.Errorf("no input messages queued for the initial turn")
 		}
-		e.mu.Lock()
-		e.started = true
-		e.mu.Unlock()
-	} else if len(input) == 0 {
 		return fmt.Errorf("Run called with no queued input and no work pending")
 	}
 
@@ -280,7 +357,9 @@ func (e *antigravityInteractionsExecution) Run(ctx context.Context, handler Hand
 		return fmt.Errorf("interaction turn failed: %w", err)
 	}
 	prevID = res.interactionID
-	e.setPrevID(prevID)
+	if err := e.setPrevID(ctx, prevID); err != nil {
+		return err
+	}
 
 	for turn := 0; turn < e.harness.cfg.MaxTurns; turn++ {
 		e.harness.debugTurn(e.conversationID, turn+1, len(res.toolCalls))
@@ -315,10 +394,12 @@ func (e *antigravityInteractionsExecution) Run(ctx context.Context, handler Hand
 
 		res, err = e.harness.postTurn(ctx, token, e.harness.newRequest(next, prevID))
 		if err != nil {
-			return fmt.Errorf("resume turn failed: %w", err)
+			return fmt.Errorf("continuation turn failed: %w", err)
 		}
 		prevID = res.interactionID
-		e.setPrevID(prevID)
+		if err := e.setPrevID(ctx, prevID); err != nil {
+			return err
+		}
 	}
 
 	// Hit the turn cap while still driving tools.
@@ -369,7 +450,7 @@ type userInputStep struct {
 	Content []textPart `json:"content"`
 }
 
-// toolResultStep is one tool result returned on a resume turn: a flat Step with
+// toolResultStep is one tool result returned on a continuation turn: a flat Step with
 // "type":"function_result", the matching call_id, the tool name, and the result
 // object under "result".
 type toolResultStep struct {
@@ -552,7 +633,7 @@ func newTokenSource(ctx context.Context) (oauth2.TokenSource, error) {
 // newRequest builds an interactionRequest common to every turn. The environment
 // is always the client-side ("local") environment -- this harness exists to
 // execute the agent's built-in env tools locally. Tools are re-declared on every
-// turn so they stay known to the agent across resumes.
+// turn so they stay known to the agent across the turns of the interaction loop.
 func (h *AntigravityInteractionsHarness) newRequest(input []any, previousID string) interactionRequest {
 	var tools []FunctionTool
 	if h.cfg.ThirdPartyExecutor != nil {
@@ -571,8 +652,93 @@ func (h *AntigravityInteractionsHarness) newRequest(input []any, previousID stri
 	}
 }
 
-// postTurn POSTs the request and streams the SSE response for one turn.
+// retryableHTTPError marks a transient HTTP response that should be retried.
+// retryAfter is the server-suggested delay (parsed from the Retry-After
+// header), or 0 if none was provided.
+type retryableHTTPError struct {
+	status     int
+	retryAfter time.Duration
+	body       string
+}
+
+func (e *retryableHTTPError) Error() string {
+	return fmt.Sprintf("HTTP %d: %s", e.status, e.body)
+}
+
+// Retry tuning for transient rate-limit (429) responses. The stateful
+// interaction quota is per-minute, so the backoff must be able to span ~a
+// minute; with these values the cumulative wait reaches ~1-2 minutes before
+// giving up.
+const (
+	rateLimitMaxRetries = 6
+	rateLimitBaseDelay  = 2 * time.Second
+	rateLimitMaxDelay   = 32 * time.Second
+)
+
+// postTurn POSTs one turn, retrying on HTTP 429 (rate limit / quota) with
+// exponential backoff and jitter, honoring a Retry-After header when present.
+// Creating an interaction is not idempotent, so only 429 -- which is rejected
+// before any interaction is created -- is retried here.
 func (h *AntigravityInteractionsHarness) postTurn(ctx context.Context, token string, reqBody interactionRequest) (*turnResult, error) {
+	delay := rateLimitBaseDelay
+	for attempt := 0; ; attempt++ {
+		res, err := h.postTurnOnce(ctx, token, reqBody)
+		if err == nil {
+			return res, nil
+		}
+		var re *retryableHTTPError
+		if !errors.As(err, &re) || attempt >= rateLimitMaxRetries {
+			return nil, err
+		}
+
+		wait := re.retryAfter
+		if wait <= 0 {
+			wait = backoffWithJitter(delay)
+		}
+		fmt.Fprintf(os.Stderr, "[harness] rate limited (HTTP %d); retrying in %s (attempt %d/%d)\n",
+			re.status, wait.Round(time.Millisecond), attempt+1, rateLimitMaxRetries)
+
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(wait):
+		}
+		if delay *= 2; delay > rateLimitMaxDelay {
+			delay = rateLimitMaxDelay
+		}
+	}
+}
+
+// backoffWithJitter returns a randomized delay in [delay/2, delay] to avoid
+// synchronized retries.
+func backoffWithJitter(delay time.Duration) time.Duration {
+	half := delay / 2
+	return half + time.Duration(rand.Int63n(int64(half)+1))
+}
+
+// parseRetryAfter parses a Retry-After header value, which may be either an
+// integer number of seconds or an HTTP date. Returns 0 if absent or unparseable.
+func parseRetryAfter(v string) time.Duration {
+	v = strings.TrimSpace(v)
+	if v == "" {
+		return 0
+	}
+	if secs, err := strconv.Atoi(v); err == nil {
+		if secs <= 0 {
+			return 0
+		}
+		return time.Duration(secs) * time.Second
+	}
+	if t, err := http.ParseTime(v); err == nil {
+		if d := time.Until(t); d > 0 {
+			return d
+		}
+	}
+	return 0
+}
+
+// postTurnOnce POSTs the request and streams the SSE response for one turn.
+func (h *AntigravityInteractionsHarness) postTurnOnce(ctx context.Context, token string, reqBody interactionRequest) (*turnResult, error) {
 	body, err := json.Marshal(reqBody)
 	if err != nil {
 		return nil, err
@@ -595,7 +761,17 @@ func (h *AntigravityInteractionsHarness) postTurn(ctx context.Context, token str
 		if _, err := b.ReadFrom(resp.Body); err != nil {
 			return nil, fmt.Errorf("HTTP %d (failed to read error body: %v)", resp.StatusCode, err)
 		}
-		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, b.String())
+		msg := b.String()
+		// 429 means the request was rejected by quota before any interaction was
+		// created, so it is safe to retry (unlike a partial 5xx). Mark it.
+		if resp.StatusCode == http.StatusTooManyRequests {
+			return nil, &retryableHTTPError{
+				status:     resp.StatusCode,
+				retryAfter: parseRetryAfter(resp.Header.Get("Retry-After")),
+				body:       msg,
+			}
+		}
+		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, msg)
 	}
 	return h.parseStreamedTurn(resp.Body)
 }

internal/harness/antigravityinteractions_tools.go

@@ -29,13 +29,14 @@ import (
 // send back as a function_result step. The built-in environment tools enumerated
 // in the switch below (file reads, command execution, and the file-mutation
 // family) are executed internally against the local filesystem/shell; every
-// other tool is dispatched to the configured ThirdPartyExecutor. All execution
-// is internal to the harness -- no tool call is surfaced to the caller.
+// other tool is dispatched to the configured ThirdPartyExecutor, or, if none is
+// configured, returned as an error result. All execution is internal to the
+// harness -- no tool call is surfaced to the caller.
 //
 // Argument names match the agent's tool schema (PascalCase). Mutation tools
-// (move/delete_dir/file_change family) require no success payload -- on success
-// they return an empty result; on failure they return {"error": <message>},
-// which marks the step as failed.
+// (move, delete_dir, create_file, edit_file, multi_edit_file, delete_file)
+// require no success payload -- on success they return an empty result; on
+// failure they return {"error": <message>}, which marks the step as failed.
 func (h *AntigravityInteractionsHarness) executeTool(ctx context.Context, call capturedToolCall) any {
 	switch call.name {
 	case "view_file":
@@ -70,7 +71,7 @@ func (h *AntigravityInteractionsHarness) executeTool(ctx context.Context, call c
 // These run against the actual local filesystem/shell because they ARE the
 // client-side environment. Argument names follow the agent's tool schema
 // (view_file -> AbsolutePath, run_command -> CommandLine, list_dir ->
-// DirectoryPath), and result field names follow what the proxy maps back into
+// DirectoryPath), and result field names follow what the server maps back into
 // the step's own output (content, Output/ExitCode, results).
 // ---------------------------------------------------------------------------