feat: complete Windows client control and recovery paths

This commit is contained in:
2026-09-06 13:37:32 +00:00
parent 56b15c7f4f
commit 486894557d
38 changed files with 2188 additions and 106 deletions
+17 -6
View File
@@ -64,18 +64,30 @@ func (server *AgentServer) ServeHTTP(response http.ResponseWriter, request *http
http.Error(response, "agent session registry is unavailable", http.StatusServiceUnavailable)
return
}
started := time.Now()
heartbeat := newSynchronizedHeartbeat(server.heartbeatIdle(), server.livenessTimeout(), 0)
connection, err := websocket.Accept(response, request, &websocket.AcceptOptions{
CompressionMode: websocket.CompressionDisabled,
// coder/websocket consumes control frames inside Read. Count ping and
// pong callbacks as inbound activity so a healthy, otherwise idle agent
// is not mistaken for a dead peer; application Read results are observed
// by serveConnection as usual.
OnPingReceived: func(context.Context, []byte) bool {
heartbeat.Observe(time.Since(started))
return true
},
OnPongReceived: func(context.Context, []byte) {
heartbeat.Observe(time.Since(started))
},
})
if err != nil {
return
}
defer connection.CloseNow()
server.serveConnection(request.Context(), connection)
server.serveConnection(request.Context(), connection, heartbeat, started)
}
func (server *AgentServer) serveConnection(parent context.Context, connection *websocket.Conn) {
started := time.Now()
func (server *AgentServer) serveConnection(parent context.Context, connection *websocket.Conn, heartbeat *synchronizedHeartbeat, started time.Time) {
messageType, payload, err := connection.Read(parent)
if err != nil {
return
@@ -194,7 +206,6 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
return
}
heartbeat := newSynchronizedHeartbeat(server.heartbeatIdle(), server.livenessTimeout(), 0)
writerDone := make(chan struct{})
go func() {
defer close(writerDone)
@@ -266,7 +277,7 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
server.close(connection, websocket.StatusPolicyViolation, "invalid command acceptance")
return
}
if _, acceptErr := server.Store.RecordCommandAcceptance(sessionContext, issue, hello.GetClientId(), registration.Generation, acknowledgement.GetCommandRevision(), acknowledgement.GetAccepted(), server.now()); acceptErr != nil {
if _, acceptErr := server.Store.RecordCommandAcceptanceWithRejection(sessionContext, issue, hello.GetClientId(), registration.Generation, acknowledgement.GetCommandRevision(), acknowledgement.GetAccepted(), acknowledgement.GetRejection(), server.now()); acceptErr != nil {
server.close(connection, websocket.StatusPolicyViolation, "invalid command acceptance")
return
}
@@ -360,6 +371,7 @@ func eventAppendFromWire(event *rvboxv1.CommandEvent, clientID string, generatio
if lifecycle := event.GetLifecycle(); lifecycle != nil {
value := lifecycle.GetLifecycle()
result.Lifecycle, result.LifecycleRevision = &value, lifecycle.GetCommandRevision()
result.WindowsIdentity = lifecycle.GetWindowsExecutionIdentity()
}
if output := event.GetOutput(); output != nil {
result.Stream = uint16(output.GetStream())
@@ -707,7 +719,6 @@ func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.
if err != nil {
return
}
heartbeat.Observe(time.Since(started))
case HeartbeatClose:
return
}