From 1fb9bd827f6f3203d733a0bac6a7eccffb0ff261 Mon Sep 17 00:00:00 2001 From: mrhid6 Date: Fri, 31 Jul 2026 17:10:59 +0100 Subject: [PATCH] feat: Added ping command --- agent/internal/grpc/pb/vantage.pb.go | 6 +++ agent/internal/sync/sync.go | 70 ++++++++++++++++++++++++++- claude.md | 13 ++++- proto/vantage/v1/vantage.proto | 13 +++++ server/internal/grpc/pb/vantage.pb.go | 6 +++ server/internal/grpc/server.go | 21 ++++++++ 6 files changed, 127 insertions(+), 2 deletions(-) diff --git a/agent/internal/grpc/pb/vantage.pb.go b/agent/internal/grpc/pb/vantage.pb.go index d96cd29..94cec48 100644 --- a/agent/internal/grpc/pb/vantage.pb.go +++ b/agent/internal/grpc/pb/vantage.pb.go @@ -166,8 +166,14 @@ type ServerCommand struct { RunStep *RunStepCmd `json:"run_step,omitempty"` CleanupWorkspace *CleanupWorkspaceCmd `json:"cleanup_workspace,omitempty"` OpenProxy *OpenProxyCmd `json:"open_proxy,omitempty"` + Ping *PingCmd `json:"ping,omitempty"` } +// PingCmd is a server-originated liveness beat. It carries nothing and expects +// no reply: its arrival is the entire message. See the .proto for why gRPC +// keepalive is not sufficient on its own. +type PingCmd struct{} + type CleanupWorkspaceCmd struct { diff --git a/agent/internal/sync/sync.go b/agent/internal/sync/sync.go index de06199..9ba6001 100644 --- a/agent/internal/sync/sync.go +++ b/agent/internal/sync/sync.go @@ -123,6 +123,15 @@ func poll(client *grpcclient.Client, cfg *config.Config, version string) error { // continuation of a run of failures. const streamHealthyAfter = time.Minute +// Stream staleness. The server beats every 20s, so 70s tolerates three missed +// beats before the stream is written off — high enough that a slow network or a +// briefly busy server does not cost a reconnect, low enough that an agent is +// not uncommandable for minutes after a control-plane restart. +const ( + streamStaleAfter = 70 * time.Second + streamStaleCheck = 10 * time.Second +) + func runCommandStream(ctx context.Context, cfg *config.Config) { backoff := time.Second @@ -185,7 +194,13 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error { } defer client.Close() - stream, err := client.CommandStream(ctx) + // Cancelling this context is what unblocks Recv when the stream has gone + // quiet. Without it the watchdog below would have no way to interrupt a + // read that is never going to return. + streamCtx, abandon := context.WithCancel(ctx) + defer abandon() + + stream, err := client.CommandStream(streamCtx) if err != nil { return fmt.Errorf("open stream: %w", err) } @@ -207,11 +222,64 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error { return stream.Send(msg) } + // Stream liveness, tracked here rather than left to gRPC keepalive. + // + // Keepalive operates on the transport, and behind an L7 proxy the transport + // ends at the proxy: it answers pings whether or not the server behind it + // is still running. A control-plane pod that dies therefore leaves this + // agent blocked in Recv on a stream that will never deliver another message + // and never error, while the control plane dispatches commands into it and + // the operator watches nothing happen. + // + // The watchdog only arms once a ping has actually been seen. A server too + // old to send them must not be treated as dead — that would put the agent + // in a reconnect loop against a control plane that is working perfectly. + var ( + lastMu sync.Mutex + lastRecv = time.Now() + pinged bool + ) + markRecv := func(isPing bool) { + lastMu.Lock() + lastRecv = time.Now() + if isPing { + pinged = true + } + lastMu.Unlock() + } + + go func() { + t := time.NewTicker(streamStaleCheck) + defer t.Stop() + for { + select { + case <-streamCtx.Done(): + return + case <-t.C: + lastMu.Lock() + idle, armed := time.Since(lastRecv), pinged + lastMu.Unlock() + if armed && idle > streamStaleAfter { + log.Printf("command stream silent for %s, assuming it is dead", idle.Truncate(time.Second)) + abandon() + return + } + } + } + }() + for { cmd, err := stream.Recv() if err != nil { return fmt.Errorf("recv: %w", err) } + markRecv(cmd.Ping != nil) + + // Pings carry nothing and are not acknowledged; being received is their + // whole purpose. + if cmd.Ping != nil { + continue + } if cmd.GenerateKey != nil { go handleGenerateKey(cfg, cmd) diff --git a/claude.md b/claude.md index 9a1680a..451ebd0 100644 --- a/claude.md +++ b/claude.md @@ -395,7 +395,18 @@ service Vantage { `CommandStream` is the only streaming RPC: the agent authenticates once with `AgentReady`, then the server pushes `ServerCommand`s and the agent replies with `CommandResult`, `StepResult`, or `StepOutputChunk`. -`ServerCommand` variants: `GenerateKeyCmd`, `DeleteKeyCmd`, `UpdateAgentCmd`, `ApplyUpdatesCmd`, `RunStepCmd`, `CleanupWorkspaceCmd`. +`ServerCommand` variants: `GenerateKeyCmd`, `DeleteKeyCmd`, `UpdateAgentCmd`, `ApplyUpdatesCmd`, `RunStepCmd`, `CleanupWorkspaceCmd`, `OpenProxyCmd`, `PingCmd`. + +**`PingCmd` is a liveness beat, and it is not redundant with gRPC keepalive.** +The server sends one every 20s on an otherwise idle command stream; the agent +treats 70s of silence as a dead stream and reconnects. Keepalive cannot do this +job behind an L7 proxy: the agent's HTTP/2 connection terminates at the proxy, +which answers pings on its own behalf, so a control-plane pod that dies leaves +the agent blocked in `Recv` on a stream that never delivers another message and +never errors — commands dispatched into it are silently lost while `SyncKeys` +keeps succeeding and the fleet list still shows the server `active`. The agent's +watchdog arms only **after** it has seen a first ping, so an older server that +sends none is treated as working rather than put into a reconnect loop. Key-state polling stays on the 30s `SyncKeys` interval. Full message definitions live in `proto/vantage/v1/vantage.proto`. diff --git a/proto/vantage/v1/vantage.proto b/proto/vantage/v1/vantage.proto index 5754a39..9d0b0d2 100644 --- a/proto/vantage/v1/vantage.proto +++ b/proto/vantage/v1/vantage.proto @@ -182,9 +182,22 @@ message ServerCommand { RunStepCmd run_step = 6; CleanupWorkspaceCmd cleanup_workspace = 7; OpenProxyCmd open_proxy = 8; + PingCmd ping = 9; } } +// PingCmd is a liveness beat, carrying nothing and requiring no reply. +// +// It exists because gRPC keepalive cannot prove what the agent needs to know. +// Behind an L7 proxy the agent's HTTP/2 connection terminates at the proxy, so +// keepalive pings are answered by the proxy whether or not the server behind it +// is still there. A pod that dies leaves the agent blocked in Recv on a stream +// that will never produce another message and never error — commands are +// dispatched into it and silently lost. Only traffic that originates at the +// server itself distinguishes a live stream from an orphaned one. +message PingCmd { +} + // CleanupWorkspaceCmd tells the agent to recursively remove the run's working // directory once all steps on that server have finished. message CleanupWorkspaceCmd { diff --git a/server/internal/grpc/pb/vantage.pb.go b/server/internal/grpc/pb/vantage.pb.go index 648f9c2..7237d29 100644 --- a/server/internal/grpc/pb/vantage.pb.go +++ b/server/internal/grpc/pb/vantage.pb.go @@ -158,8 +158,14 @@ type ServerCommand struct { RunStep *RunStepCmd `json:"run_step,omitempty"` CleanupWorkspace *CleanupWorkspaceCmd `json:"cleanup_workspace,omitempty"` OpenProxy *OpenProxyCmd `json:"open_proxy,omitempty"` + Ping *PingCmd `json:"ping,omitempty"` } +// PingCmd is a server-originated liveness beat. It carries nothing and expects +// no reply: its arrival is the entire message. See the .proto for why gRPC +// keepalive is not sufficient on its own. +type PingCmd struct{} + type CleanupWorkspaceCmd struct { WorkspaceId string `json:"workspace_id"` } diff --git a/server/internal/grpc/server.go b/server/internal/grpc/server.go index 8a3d123..23c29f1 100644 --- a/server/internal/grpc/server.go +++ b/server/internal/grpc/server.go @@ -212,11 +212,27 @@ func (s *vantageServer) CommandStream(stream pb.Vantage_CommandStreamServer) err } }() + // The heartbeat is what lets the agent tell a live stream from an orphaned + // one. gRPC keepalive cannot: behind an L7 proxy the agent's connection + // terminates at the proxy, which answers pings on its own behalf, so a dead + // pod leaves the agent blocked in Recv forever with commands vanishing into + // a stream nobody is serving. A message that originates here is the only + // thing that proves this process is still on the other end. + ping := time.NewTicker(pingInterval) + defer ping.Stop() + ctx := stream.Context() for { select { case <-ctx.Done(): return nil + case <-ping.C: + // A failed send is the point: it is how this side learns the stream + // is gone, which runs the deferred release and frees the agent's + // presence claim for whichever pod it reconnects to. + if err := stream.Send(&pb.ServerCommand{Ping: &pb.PingCmd{}}); err != nil { + return err + } case cmd, ok := <-ch: if !ok { return nil @@ -228,6 +244,11 @@ func (s *vantageServer) CommandStream(stream pb.Vantage_CommandStreamServer) err } } +// How often the server beats on an idle command stream. Comfortably under the +// agent's staleness threshold, so a single dropped beat does not cost a +// reconnect. +const pingInterval = 20 * time.Second + // StartGRPC serves the agent API until stop is called. // // It returns a stop function rather than serving forever because an abrupt exit