[adalundhe]

Summary

A single-replica controller restart leaves the controller unreachable for 30–45 s. The Go runtime's TaskStore client gave each call four attempts within about 0.7 s, so the turn's final save failed. The task then stayed WORKING, and the Session refused new input until someone cancelled it.

A deployed runtime reaches the controller through Substrate's egress proxy, AgentGateway. During a restart the proxy stays reachable and answers each call at once with Unavailable (upstream call failed: Connect: Connection refused). The runtime's connection stays READY, so the outage never appears as a broken connection.

The Go runtime's TaskStore client now:

• retries Unavailable and DeadlineExceeded within a two-minute budget per operation, for saves, reads and settlement;
• gives each attempt the client's 30 s call timeout, within what remains of the budget;
• waits 100 ms, doubling up to 1 s, jittered between half and all of the delay;
• keeps grpc.WaitForReady(true) on the five TaskStore calls, so a broken connection waits for reconnection instead of failing fast;
• replays the identical request on every attempt, so a replay after a lost response returns the committed version;
• returns the last RPC error and logs taskstore retries exhausted once when the budget runs out. An attempt that the budget cuts short does not replace that error. A caller that cancels or runs out of time gets its own error, without that log.

Settlement no longer shares a 30 s cleanup deadline. Each telemetry flush around it keeps its own telemetry.FlushTimeout. The change adds no configuration, and Memory calls keep failing fast.

How the design changed

The first revision added a configurable budget. Review asked whether WaitForReady alone could cover a restart, and an earlier head of this PR switched to it. A kind run showed that it cannot, because the proxy keeps the connection READY. This head follows the review suggestion to keep WaitForReady and bound the existing loop by elapsed time, without new settings.

Where the guarantee stops

• A controller outage longer than two minutes still fails the save and leaves the task WORKING until someone cancels it.
• The Python runtime's TaskStore client keeps its fixed four attempts.

Validation

• go test -race ./adk/pkg/taskstore ./adk/pkg/controllerclient passes three runs. The new tests:
  • reject more calls than the former limit through a reachable fake proxy, for a save and for settlement;
  • check the budget's returned error and single log record, an attempt the budget cuts short, the caller's own cancellation and deadline, and the jitter bounds;
  • pass 20 runs, and 10 more with the CPU saturated.
• TestRuntimeTaskStoreSaveOutlastsAControllerRestart runs the controller's gRPC server on PostgreSQL in two variants: direct, and through a forwarding proxy that answers Unavailable while the controller is down.
  • Both variants pass on Linux.
  • With the previous store.go, the proxy variant fails within 0.75 s.
• make lint reports 0 issues.
• A kind run with Substrate v1.0.0-rc1 and AgentGateway v1.7.0 released a turn's final save into a 30 s controller outage. Every arm used the same controller; only the runtime image differed.

Runtime	Final save during the outage	Task afterwards	Next message
main	4 attempts in 713 ms	WORKING	refused, FailedPrecondition
7f2d5f83, WaitForReady only	4 attempts in 713 ms	WORKING	refused, FailedPrecondition
this head	37 rejected attempts, then lands as the controller returns after 34 s	COMPLETED	answered

🤖 Generated with Claude Code

[EItanya]

🤖 Includes AI-generated review comments.

[EItanya]

Should we passthrough the TaskStore as the Dependency here rather than just the retry logic?

[adalundhe]

This change is gone. taskstore.New reads KAGENT_TASKSTORE_OPERATION_BUDGET and KAGENT_TASKSTORE_ATTEMPT_TIMEOUT itself, so main.go matches main and a Harness sets the limits in spec.env.

[EItanya]

private test seam fields seem weird to me

[adalundhe]

Removed. The tests set the limits through the same environment variables an operator uses.

[EItanya]

What's the fundamental difference between taskstore/retry.go and a2a/taskstore_retry.go. Feels like a bunch of potential duplication?

[adalundhe]

api/a2a/taskstore_retry.go went away with the controller-side configuration. taskstore/retry.go is the only copy.

[EItanya]

This feels like the type of test i'd rather run as some sort of fancy integration where we run the server and client rather than FULL e2e, it's bound to take a long time and add unecessary complexity

[adalundhe]

Replaced with TestRuntimeTaskStoreSaveOutlastsAControllerRestart in core/internal/grpcserver. It runs the real controller gRPC server on Postgres and stops it while a save is pending. It restarts the server 2 s later and requires the save to land once, with the committed version. It takes about 25 s and needs only Docker.

[EItanya]

Can we give each flush its own timeout here? AppConfig.Flush accepts arbitrary callbacks, and context.WithoutCancel(ctx) removes the deadline the callback previously received. A custom context-aware flush that relies on that deadline can now block before SettleTask indefinitely, leaving the task unpublished. The built-in telemetry provider supplies its own timeout, so that path is covered. Separate flush deadlines would keep this bounded without consuming settlement's retry budget.

[adalundhe]

Done. Each flush gets its own telemetry.FlushTimeout deadline, separate from settlement's budget. In a new test, a flush that blocks on its context ends at that deadline, and settlement still retries for its whole budget.

[EItanya]

This can return context.DeadlineExceeded while the caller context is still live. For example, with an 80 ms caller deadline, an immediate Unavailable response returns a timeout after only a few milliseconds because the next backoff won't fit. That reports controller unavailability as a caller timeout. Can we wait for the caller deadline to actually expire, or preserve the last RPC error if we choose to stop early?

[adalundhe]

Done. The loop keeps waiting when the deadline falls inside a backoff, so a caller gets DeadlineExceeded only after its deadline has passed. Tests cover 80 ms and 200 ms deadlines against an immediate Unavailable. When the budget runs out first, the caller still gets the last RPC error.

[EItanya]

Optional: can we pass the retry policy when constructing the store instead of adding a setter? Every current caller sets it immediately after New, so mutable configuration adds an extra initialization step without a current need.

[adalundhe]

Done. There is no setter. New reads the limits when it builds the store.

[adalundhe]

Before and after: a turn's final save through a 30 s controller outage

I replayed the controller-restart scenario on a local kind cluster with gVisor workers. Both runs use one controller and UI, built from this PR's head (f2b34211). Only the kagent Harness's Go ADK image differs: main at a77e158b before, this PR after. Neither run sets KAGENT_TASKSTORE_*, so the after run uses the defaults.

A scheduled run sends "What is 2+2?", and a mock model holds the turn. The controller scales to zero. Once its Pod is gone, the model answers, so the turn's final save meets no controller. After 30 s the controller scales back to one replica. The driver then watches for 3 minutes and sends a follow-up on the same session.

	before (main)	after (this PR)
Controller unavailable	33 s	35 s
Runtime's final save	fails at 17:53:35 with Unavailable … Connection refused, 20 s before the controller returns	retried 25 times over 33.1 s; lands 0.1 s after the controller is Ready
Scheduled-run execution	still RUNNING after 3 minutes	SUCCEEDED within 1 s of the controller being Ready
Task	WORKING, no answer	COMPLETED: "The answer is 4."
Follow-up on the same session	refused: FailedPrecondition "resource conflict"	accepted and answered
Model calls for the turn	1	1

before	after

The runtime's own log for each turn:

before 17:53:35.233 WARN agent execution failed
       cause: processor error: failed to store failed task state: failed to save task state: previous task
       persistence failed: failed to save task state: rpc error: code = Unavailable desc = upstream call
       failed: Connect: Connection refused (os error 111)
after  18:07:01.421 INFO task store call succeeded after retries
       operation=UpdateTask attempts=25 elapsed=33.13s

The restarted controller agrees. Before, it served no save for the stuck turn; that session's only call was the refused follow-up. After, it served the retried UpdateTask at 18:07:01.420 and SettleTask at 18:07:01.456.

Screenshots, JSON reports, full logs and the driver are on evidence/3391-controller-restart.

[adalundhe]

Silly. Fixing.

[EItanya]

Why are these added as env vars and not config vars?

[adalundhe]

Why are these added as env vars and not config vars?

Fixed! Good callout. Thank you.

[adalundhe]

Moved to config in 7a2638f, following #3434; the env vars are gone.

• runtimeDefaults.taskStore (Helm: config.runtimeDefaults.taskStore) holds operationBudget and attemptTimeout. The schema bounds both and rejects an attempt longer than the budget.
• The controller compiles it into the Go ADK runtime's runtime.task_store only when one is set, so an older runtime image still loads its configuration.
• The Go ADK main passes it as app.AppConfig.TaskStore, and taskstore.New takes it. Absent fields keep the 90 s and 30 s defaults.

Claude and Codex keep the defaults for now.

[jjamroga]

Could gRPC's WaitForReady cover most of this, instead of a custom retry loop and new settings?

I tried a smaller change on main. It adds grpc.WaitForReady(true) to the five TaskStore calls (Create, Update, Get, List, Settle) and leaves the existing 4-attempt loop and 30 s client timeout as they are. During a restart, a call then waits for the channel to reconnect instead of failing fast with Unavailable. The deadline still bounds every call.

I ran this PR's TestRuntimeTaskStoreSaveOutlastsAControllerRestart against it, with the outage length as a parameter. The only test change is the assertion that a save made more than one attempt, since a save may now wait in a single call:

Controller down	Result	UpdateTask attempts
2 s (unmodified main)	fails as expected	—
2 s	pass	1
10 s	pass	1
45 s	pass	2
100 s	pass	4

4 attempts × 30 s covers about 2 minutes, more than the 90 s default here. I scoped WaitForReady to the TaskStore calls so memory calls still fail fast.

If this direction works, would it make sense to keep only these parts of this PR?

• Remove settlement's 30 s cap in Cleanup and keep the per-flush telemetry.FlushTimeout.
• Keep the log line when retries are exhausted.

Then retry.go, the TaskStoreConfig proto and its validation, the Helm value and the translator changes could go. A configurable budget could come later if we need one.

Things I haven't covered:

• a call already in flight when the controller shuts down (the existing loop should retry its Unavailable);
• a kind run;
• Python. _task_store.py still uses a fixed 4 attempts and would need wait_for_ready=True too.

Is there a case WaitForReady misses that needs the budget loop?

[adalundhe]

Switched this PR to the scoped WaitForReady proposal in 7f2d5f83.

All five TaskStore RPCs now use grpc.WaitForReady(true), retaining four attempts and the 30-second per-call deadline. Settlement drops its shared deadline; each telemetry flush gets telemetry.FlushTimeout. Retry exhaustion emits an error log. The custom retry budget and its proto, Helm and translator changes are gone.

The evidence bundle contains raw logs, exported spans, hashes, and a rerunnable driver pinned to this commit. These fresh adversarial checks use controlled gRPC services and -race:

Case	Observed result
Baseline 12aa4d46, unavailable transport	Expected failure: a TaskStore RPC returns Unavailable before recovery
Head, 100ms unavailable transport	All five TaskStore calls succeed in 103.4–103.7ms; Memory fails fast in 0.139ms
Caller cancellation / earlier deadline	Stop in 259ms / 260ms, one attempt each
Permanent transport outage	Four 30-second attempts; DeadlineExceeded after 120.701851s; exhaustion log observed
RPC already in flight when transport stops	2.002107s outage; two identical requests; replacement service acknowledges version 2 after 2.857552s total
READY server returns Unavailable	Four attempts exhaust after 728.646ms; synthetic backend recovers after 2.002615s
Canceled observer plus expired telemetry flush	Regression passes: settlement and the second flush get fresh deadlines

Two replay runs failed timing assertions in the evidence driver: exhaustion returned after 124.071344s, and cancellation/deadline results arrived after 851ms. Attempt counts and error codes remained correct. Both failures remain in the bundle. The revised driver checks configured deadlines, preservation of earlier caller deadlines, cancellation and attempt counts. The final full replay matched all five expected outcomes, including the baseline failure; exhaustion returned after 120.702011s.

The READY-server case confirms a boundary: WaitForReady waits for channel readiness, not recovery from application-returned errors. It uses a synthetic service; it does not reproduce a deployed proxy failure or establish a need for the removed budget.

Current-head CI passed Go race tests, Go lint, and both gVisor and microVM E2E. CI tested merge 37d67349, combining this head with main 82bddc8d. The Go race job includes the controller gRPC package containing the PostgreSQL restart regression.

The bundle retains six fresh OTel client spans and an allowlisted extract of current CI traces and actor logs. The original CI artifact includes controller, worker, Substrate and test logs; it expires November 9. Span IDs, timestamps and RPC statuses are retained. The selected CI trace's association with the scheduled-restart test is a timing inference, stated in the report.

Fresh deployed final-save recovery during an outage remains unverified. Docker Desktop's required sign-in blocked local PostgreSQL and Kind runs before container creation. The prepared 2/10/45/100-second controller matrix did not run. CI's scheduled-restart test releases the model after controller readiness, so it does not force the final save into the outage. The earlier before/after deployment evidence covers f2b34211 and its custom budget; it is historical evidence rather than proof of this switch.

[EItanya]

🤖 AI-generated review.

The READY-server case in the latest evidence is the remaining concern. When the runtime's connection to AgentGateway stays healthy but the controller upstream is down, the proxy can return Unavailable immediately. WaitForReady doesn't wait for recovery in that case, so four attempts still exhaust in under a second and leave the execution unable to save after recovery. The earlier deployed run showed this upstream-connect-refused error.

Can we keep WaitForReady and bound the existing retry loop by elapsed time? This is the small patch tested locally against 7f2d5f83. It keeps the 30-second per-call timeout, uses a two-minute total budget, and caps backoff at one second without adding configuration:

 func (s *Store) retry(ctx context.Context, call func(context.Context) error) error {
-	for attempt := range 4 {
+	// Bound the whole operation, including immediate failures returned by a
+	// healthy proxy while its upstream is down. Keep the former four 30-second
+	// attempts' overall allowance without limiting the number of fast failures.
+	ctx, cancelOperation := context.WithTimeout(ctx, 2*time.Minute)
+	defer cancelOperation()
+	delay := 100 * time.Millisecond
+	for attempt := 1; ; attempt++ {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
 		rpcCtx, cancel := s.callContext(ctx)
 		err := call(rpcCtx)
 		cancel()
 		if status.Code(err) != codes.Unavailable && status.Code(err) != codes.DeadlineExceeded {
 			return err
 		}
-		if attempt == 3 {
-			logging.FromContext(ctx).ErrorContext(ctx, "taskstore retries exhausted", "attempts", attempt+1, "error", err)
-			return err
-		}
-		timer := time.NewTimer(time.Duration(1<<attempt) * 100 * time.Millisecond)
+		timer := time.NewTimer(delay)
 		select {
 		case <-ctx.Done():
 			timer.Stop()
+			logging.FromContext(ctx).ErrorContext(ctx, "taskstore retries exhausted", "attempts", attempt, "error", err)
 			return ctx.Err()
 		case <-timer.C:
 		}
+		delay = min(2*delay, time.Second)
 	}
-	panic("unreachable")
 }

With a two-second simulated upstream outage over a healthy gRPC connection, the current head failed 3/3 runs after about 704 ms. With this patch, all three runs recovered after about 2.5 seconds and six attempts, and subsequent saves in the same execution succeeded. Caller deadline/cancellation and terminal-error checks, the TaskStore/controller-client tests, and the PostgreSQL controller-restart test also passed under -race.

This was a local test using real gRPC with simulated proxy responses; it wasn't a live AgentGateway/Kind run. The regression should also cover a reachable proxy returning Unavailable while the controller is down.

[adalundhe]

WaitForReady alone misses the deployed case, and the suggested elapsed-time loop covers it. 8f893624 applies that loop with three changes, and it passed the adversarial pass below. Our earlier evidence comment was wrong to say the READY-server case did not reproduce a deployed failure.

A deployed runtime reaches the controller through the egress proxy, AgentGateway. During a restart the proxy answers each TaskStore call at once with Unavailable (upstream call failed: Connect: Connection refused (os error 111)). The runtime's connection to it stays up, so WaitForReady never waits. Four attempts ran out in about 0.7 s.

8f893624 bounds each operation by a two-minute budget and keeps WaitForReady, the 30 s call timeout, and 100 ms–1 s waits. Its changes to the suggested loop:

• It logs taskstore retries exhausted only when the budget runs out, not when the caller cancels or its own deadline passes.
• On exhaustion it returns the last RPC error, so the runtime's failure keeps the proxy's cause instead of context deadline exceeded.
• It jitters each wait between half and all of the delay.

The evidence bundle has the logs, traces, screenshots, drivers and checksums.

Live kind runs

A local kind cluster with gVisor workers, Substrate v1.0.0-rc1, AgentGateway v1.7.0 and an OpenTelemetry collector ran both arms. The arms shared the cluster and controller; only the runtime image differed. The driver released a turn's final save into a 30 s controller outage.

	7f2d5f83 (WaitForReady, 4 attempts)	8f893624
Controller unavailable	35 s	35 s
Final save	4 UNAVAILABLE spans in 717 ms, then taskstore retries exhausted attempts=4	38 UNAVAILABLE spans, then UpdateTask OK and SettleTask OK in the same trace
Task	WORKING after 3 minutes	COMPLETED 1 s after the controller was Ready
Follow-up on the same Session	refused, FailedPrecondition "resource conflict"	answered

All failed attempts in each arm went over one runtime connection to the proxy.

Screenshots: chat after each run

7f2d5f83: the prompt never gets an answer.

8f893624: the turn and the follow-up both complete.

Traces: the runtime's TaskStore client spans

7f2d5f83, trace 748409c41d1666db:

8f893624, trace ff80d22e22358662:

Adversarial pass

Probes at 8f893624 run real gRPC over bufconn under -race. A fake server answers like the proxy during a restart, over a connection that stays READY:

Case	Result
2 s outage behind a reachable proxy	lands after 2.1 s and 6 identical attempts, all on a READY channel
Settlement through the same outage	settles after 2.3 s
Proxy unreachable for 2 s, then the controller down for 2 s	lands after 4.1 s
Controller never returns	stops after 2m0.002s and 161 attempts; returns the proxy's Unavailable and logs once
Caller cancels at 250 ms	returns context canceled; logs nothing
Caller deadline at 1.5 s	returns context deadline exceeded at 1.5 s; logs nothing

An independent review of the change found one bug, which 8f893624 fixes. When the budget ended during an attempt, that attempt's own deadline error replaced the proxy's error. The context deadline can also fire late under load, so one more attempt could start after the budget had passed. The loop now keeps the error from before a cut-short attempt and checks the clock before each attempt. With a proxy taking 300 ms per attempt, the loop before the fix returned DeadlineExceeded in 10 of 10 runs, and 8f893624 returned the proxy's Unavailable in 10 of 10. The new regression test fails 5 of 5 runs against the loop before the fix.

The PR's unit tests cover the proxy outage, settlement, exhaustion, the caller's own end, a cut-short attempt and the jitter bounds. They pass 20 runs, and 10 more with every CPU core saturated.

Controller and PostgreSQL

TestRuntimeTaskStoreSaveOutlastsAControllerRestart now runs two variants: direct, and through a forwarding proxy that answers Unavailable while the controller is down. Both pass on Linux. With 7f2d5f83's store.go, the proxy variant fails within 0.76 s. make lint reports 0 issues, and CI Build passes at 8f893624.

Where the evidence stops

• Each live arm ran once, with one 30 s outage.
• An outage longer than two minutes still fails the save and leaves the task WORKING.
• The Python runtime's TaskStore client keeps four attempts.