// Add to coderd/workspaceagents.go

// authenticateAgent handles workspace agent authentication, with support for restored agents
func (api *API) authenticateAgent(r *http.Request) (*database.WorkspaceAgent, error) {
    token := r.Header.Get(AgentAuthTokenHeader)
    if token == "" {
        return nil, xerrors.New("no agent auth token provided")
    }

    agent, err := api.Database.GetWorkspaceAgentByAuthToken(r.Context(), token)
    if err != nil {
        if !xerrors.Is(err, sql.ErrNoRows) {
            return nil, xerrors.Errorf("get agent by auth token: %w", err)
        }
        
        // Check if this is a restored agent trying to reconnect
        restoredAgent, restoredErr := api.tryAuthenticateRestoredAgent(r.Context(), token)
        if restoredErr == nil && restoredAgent != nil {
            // Found a restored agent, regenerate auth token and return
            return restoredAgent, nil
        }
        
        return nil, xerrors.Errorf("agent auth token invalid: %w", err)
    }

    return agent, nil
}

// tryAuthenticateRestoredAgent attempts to authenticate an agent that may have been restored from backup
func (api *API) tryAuthenticateRestoredAgent(ctx context.Context, token string) (*database.WorkspaceAgent, error) {
    // Try to find an agent that matches the first part of the token (pre-secret portion)
    // This is a heuristic approach - in a real implementation you'd need a more secure method
    tokenParts := strings.Split(token, ".")
    if len(tokenParts) < 2 {
        return nil, xerrors.New("invalid token format")
    }
    
    agentID := tokenParts[0]
    
    agents, err := api.Database.GetWorkspaceAgentsCreatedAfter(ctx, time.Now().Add(-7*24*time.Hour))
    if err != nil {
        return nil, err
    }
    
    // Find a potential restored agent
    for _, agent := range agents {
        if strings.HasPrefix(agent.ID.String(), agentID) {
            // Found a potential match, regenerate auth for this agent
            newToken := uuid.New().String()
            err = api.Database.UpdateWorkspaceAgentAuthToken(ctx, database.UpdateWorkspaceAgentAuthTokenParams{
                ID: agent.ID,
                AuthToken: newToken,
            })
            if err != nil {
                return nil, err
            }
            
            // Log the recovery for audit purposes
            api.Logger.Info(ctx, "restored workspace agent authentication",
                slog.F("agent_id", agent.ID),
                slog.F("workspace_id", agent.WorkspaceID),
            )
            
            return agent, nil
        }
    }
    
    return nil, xerrors.New("no matching restored agent found")
}

// Also add needed database methods to coderd/database/database.go:

// GetWorkspaceAgentsCreatedAfter gets all workspace agents created after the given time
func (q *Q) GetWorkspaceAgentsCreatedAfter(ctx context.Context, after time.Time) ([]database.WorkspaceAgent, error) {
    // Implementation details will depend on your database schema
    // ...
}

// UpdateWorkspaceAgentAuthToken updates the auth token for a workspace agent
func (q *Q) UpdateWorkspaceAgentAuthToken(ctx context.Context, params database.UpdateWorkspaceAgentAuthTokenParams) error {
    // Implementation details will depend on your database schema
    // ...
}

// In coderd/workspaceagents.go

func (api *API) registerAgentRecoveryHandlers(r *chi.Mux) {
    r.Post("/api/v2/workspaceagents/recovery", api.handleAgentRecovery)
}

func (api *API) handleAgentRecovery(rw http.ResponseWriter, r *http.Request) {
    // Extract agent identity information from request
    var req struct {
        AgentID       string `json:"agent_id"`
        WorkspaceID   string `json:"workspace_id"`
        RecoveryToken string `json:"recovery_token"`  // A token derived from agent's instance data
    }
    
    if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
        httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
            Message: "Invalid request",
            Detail:  err.Error(),
        })
        return
    }
    
    // Validate this is a legitimate agent through instance verification 
    // ...
    
    // Regenerate token and update database
    // ...
    
    // Return new auth token to the agent
}