added the logic to populate the PeerUpdate after updating the internal tunnel state

ibetitsmike · ibetitsmike · commit f66d81a44457 · 2025-04-29T10:49:52.000Z
diff --git a/vpn/tunnel.go b/vpn/tunnel.go
@@ -397,6 +397,38 @@ func (u *updater) sendUpdateResponse(req *request[*TunnelMessage, *ManagerMessag
 // createPeerUpdateLocked creates a PeerUpdate message from a workspace update, populating
 // the network status of the agents.
 func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUpdate {
+
+	// this flag is true on the first update after a reconnect
+	if update.FreshState {
+		// ignoredWorkspaces is initially populated with the workspaces that are
+		// in the current update. Later on we populate it with the deleted workspaces too
+		// so that we don't send duplicate updates.
+		ignoredWorkspaces := make(map[uuid.UUID]struct{}, len(update.UpsertedWorkspaces))
+
+		for _, workspace := range update.UpsertedWorkspaces {
+			ignoredWorkspaces[workspace.ID] = struct{}{}
+		}
+		for _, agent := range update.UpsertedAgents {
+			// delete any current agents that are not in the new update
+			if _, ok := u.agents[agent.ID]; !ok {
+				delete(u.agents, agent.ID)
+				update.DeletedAgents = append(update.DeletedAgents, &tailnet.Agent{
+					ID:          agent.ID,
+					Name:        agent.Name,
+					WorkspaceID: agent.WorkspaceID,
+				})
+				if _, ok := ignoredWorkspaces[agent.WorkspaceID]; !ok {
+					update.DeletedWorkspaces = append(update.DeletedWorkspaces, &tailnet.Workspace{
+						// other fields cannot be populated because the tunnel
+						// only stores agents and corresponding workspaceIDs
+						ID: agent.WorkspaceID,
+					})
+					ignoredWorkspaces[agent.WorkspaceID] = struct{}{}
+				}
+			}
+		}
+	}
+
 	out := &PeerUpdate{
 		UpsertedWorkspaces: make([]*Workspace, len(update.UpsertedWorkspaces)),
 		UpsertedAgents:     make([]*Agent, len(update.UpsertedAgents)),
@@ -413,16 +445,6 @@ func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUp
 		delete(u.agents, agent.ID)
 	}
 
-	// this flag is true on the first update after a reconnect
-	if update.FreshState {
-		for _, agent := range update.UpsertedAgents {
-			// delete any current agents that are not in the new update
-			if _, ok := u.agents[agent.ID]; !ok {
-				delete(u.agents, agent.ID)
-			}
-		}
-	}
-
 	for i, ws := range update.UpsertedWorkspaces {
 		out.UpsertedWorkspaces[i] = &Workspace{
 			Id:     tailnet.UUIDToByteSlice(ws.ID),
